what.models.detection.datasets.fiftyone

 1import torch
 2import cv2
 3import numpy as np
 4
 5class FiftyOneDataset(torch.utils.data.Dataset):
 6    """A class to construct a PyTorch dataset from a FiftyOne dataset.
 7
 8    Args:
 9        fiftyone_dataset: a FiftyOne dataset or view that will be used for 
10            training or testing
11        transform (None): a list of PyTorch transform to apply to images 
12            and targets when loading
13        ground_truth_field ("ground_truth"): the name of the field in fiftyone_dataset 
14            that contains the desired labels to load
15        classes (None): a list of class strings that are used to define the 
16            mapping between class names and indices. If None, it will use 
17            all classes present in the given fiftyone_dataset.
18    """
19
20    def __init__(self, fiftyone_dataset, classes,
21                 transform = None, target_transform = None,
22                 ground_truth_field = "ground_truth"):
23
24        self.dataset = fiftyone_dataset
25        self.transform = transform
26        self.target_transform = target_transform
27        self.ground_truth_field = ground_truth_field
28
29        self.img_paths = self.dataset.values("filepath")
30
31        self.classes = classes
32
33        if self.classes[0] not in ["BACKGROUND", "background"]:
34            self.classes = ["BACKGROUND"] + self.classes
35
36        self.labels_map_rev = {c: i for i, c in enumerate(self.classes)}
37
38    def __getitem__(self, idx):
39        img_path = self.img_paths[idx]
40        sample = self.dataset[img_path]
41
42        img = cv2.imread(str(img_path))
43        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
44        height, width, _ = img.shape
45
46        boxes = []
47        labels = []
48        detections = sample[self.ground_truth_field].detections
49        for det in detections:
50            category_id = self.labels_map_rev[det.label]
51            x, y, w, h = det.bounding_box
52            boxes.append(np.array([float(x)*width, float(y)*height, float(x + w)*width, float(y + h)*height]))
53            labels.append(category_id)
54        boxes = np.array(boxes, dtype=np.float32)
55        labels = np.array(labels, dtype=np.int64)
56
57        # Background
58        if len(boxes) == 0:
59            boxes = np.array([[0, 0, img.shape[1], img.shape[0]]], dtype=np.float32)
60            labels = np.array([0], dtype=np.int64)
61    
62        if self.transform:
63            img, boxes, labels = self.transform(img, boxes, labels)
64        if self.target_transform:
65            boxes, labels = self.target_transform(boxes, labels)
66                
67        return img, boxes, labels
68
69    def __len__(self):
70        return len(self.img_paths)
71
72    def get_classes(self):
73        return self.classes