what.models.detection.datasets.open_images

  1import copy
  2import pathlib
  3import cv2
  4import numpy as np
  5import pandas as pd
  6
  7class OpenImagesDataset:
  8
  9    def __init__(self, root,
 10                 transform=None, target_transform=None,
 11                 dataset_type="train", balance_data=False):
 12        self.root = pathlib.Path(root)
 13        self.transform = transform
 14        self.target_transform = target_transform
 15        self.dataset_type = dataset_type.lower()
 16
 17        self.data, self.class_names, self.class_dict = self._read_data()
 18        self.balance_data = balance_data
 19        self.min_image_num = -1
 20        if self.balance_data:
 21            self.data = self._balance_data()
 22        self.ids = [info['image_id'] for info in self.data]
 23
 24        self.class_stat = None
 25
 26    def _getitem(self, index):
 27        image_info = self.data[index]
 28        image = self._read_image(image_info['image_id'])
 29        # duplicate boxes to prevent corruption of dataset
 30        boxes = copy.copy(image_info['boxes'])
 31        boxes[:, 0] *= image.shape[1]
 32        boxes[:, 1] *= image.shape[0]
 33        boxes[:, 2] *= image.shape[1]
 34        boxes[:, 3] *= image.shape[0]
 35        # duplicate labels to prevent corruption of dataset
 36        labels = copy.copy(image_info['labels'])
 37        if self.transform:
 38            image, boxes, labels = self.transform(image, boxes, labels)
 39        if self.target_transform:
 40            boxes, labels = self.target_transform(boxes, labels)
 41        return image_info['image_id'], image, boxes, labels
 42
 43    def __getitem__(self, index):
 44        _, image, boxes, labels = self._getitem(index)
 45        return image, boxes, labels
 46
 47    def get_annotation(self, index):
 48        """To conform the eval_ssd implementation that is based on the VOC dataset."""
 49        image_id, image, boxes, labels = self._getitem(index)
 50        is_difficult = np.zeros(boxes.shape[0], dtype=np.uint8)
 51        return image_id, (boxes, labels, is_difficult)
 52
 53    def get_image(self, index):
 54        image_info = self.data[index]
 55        image = self._read_image(image_info['image_id'])
 56        if self.transform:
 57            image, _ = self.transform(image)
 58        return image
 59
 60    def _read_data(self):
 61        annotation_file = f"{self.root}/sub-{self.dataset_type}-annotations-bbox.csv"
 62        annotations = pd.read_csv(annotation_file)
 63        class_names = ['BACKGROUND'] + sorted(list(annotations['ClassName'].unique()))
 64        class_dict = {class_name: i for i, class_name in enumerate(class_names)}
 65        data = []
 66        for image_id, group in annotations.groupby("ImageID"):
 67            boxes = group.loc[:, ["XMin", "YMin", "XMax", "YMax"]].values.astype(np.float32)
 68            # make labels 64 bits to satisfy the cross_entropy function
 69            labels = np.array([class_dict[name] for name in group["ClassName"]], dtype='int64')
 70            data.append({
 71                'image_id': image_id,
 72                'boxes': boxes,
 73                'labels': labels
 74            })
 75        return data, class_names, class_dict
 76
 77    def __len__(self):
 78        return len(self.data)
 79
 80    def __repr__(self):
 81        if self.class_stat is None:
 82            self.class_stat = {name: 0 for name in self.class_names[1:]}
 83            for example in self.data:
 84                for class_index in example['labels']:
 85                    class_name = self.class_names[class_index]
 86                    self.class_stat[class_name] += 1
 87        content = ["Dataset Summary:"
 88                   f"Number of Images: {len(self.data)}",
 89                   f"Minimum Number of Images for a Class: {self.min_image_num}",
 90                   "Label Distribution:"]
 91        for class_name, num in self.class_stat.items():
 92            content.append(f"\t{class_name}: {num}")
 93        return "\n".join(content)
 94
 95    def _read_image(self, image_id):
 96        image_file = self.root / self.dataset_type / f"{image_id}.jpg"
 97        image = cv2.imread(str(image_file))
 98        if image.shape[2] == 1:
 99            image = cv2.cvtColor(image, cv2.COLOR_GRAY2RGB)
100        else:
101            image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
102        return image
103
104    def _balance_data(self):
105        label_image_indexes = [set() for _ in range(len(self.class_names))]
106        for i, image in enumerate(self.data):
107            for label_id in image['labels']:
108                label_image_indexes[label_id].add(i)
109        label_stat = [len(s) for s in label_image_indexes]
110        self.min_image_num = min(label_stat[1:])
111        sample_image_indexes = set()
112        for image_indexes in label_image_indexes[1:]:
113            image_indexes = np.array(list(image_indexes))
114            sub = np.random.permutation(image_indexes)[:self.min_image_num]
115            sample_image_indexes.update(sub)
116        sample_data = [self.data[i] for i in sample_image_indexes]
117        return sample_data