what.models.detection.datasets.voc

  1import os
  2import pathlib
  3import numpy as np
  4import xml.etree.ElementTree as ET
  5import cv2
  6
  7VOC_CLASS_NAMES =  ["BACKGROUND",
  8                    "aeroplane",
  9                    "bicycle",
 10                    "bird",
 11                    "boat",
 12                    "bottle",
 13                    "bus",
 14                    "car",
 15                    "cat",
 16                    "chair",
 17                    "cow",
 18                    "diningtable",
 19                    "dog",
 20                    "horse",
 21                    "motorbike",
 22                    "person",
 23                    "pottedplant",
 24                    "sheep",
 25                    "sofa",
 26                    "train",
 27                    "tvmonitor"]
 28
 29class VOCDataset:
 30
 31    def __init__(self, root, transform=None, target_transform=None, is_test=False, keep_difficult=False, label_file=None):
 32        """Dataset for VOC data.
 33        Args:
 34            root: the root of the VOC2007 or VOC2012 dataset, the directory contains the following sub-directories:
 35                Annotations, ImageSets, JPEGImages, SegmentationClass, SegmentationObject.
 36        """
 37        self.root = pathlib.Path(root)
 38        self.transform = transform
 39        self.target_transform = target_transform
 40        if is_test:
 41            image_sets_file = self.root / "ImageSets/Main/test.txt"
 42        else:
 43            image_sets_file = self.root / "ImageSets/Main/trainval.txt"
 44        self.ids = VOCDataset._read_image_ids(image_sets_file)
 45        self.keep_difficult = keep_difficult
 46
 47        # if the labels file exists, read in the class names
 48        label_file_name = self.root / "labels.txt"
 49
 50        if os.path.isfile(label_file_name):
 51            class_string = ""
 52            with open(label_file_name, 'r') as infile:
 53                for line in infile:
 54                    class_string += line.rstrip()
 55
 56            # classes should be a comma separated list
 57            
 58            classes = class_string.split(',')
 59            # prepend BACKGROUND as first class
 60            classes.insert(0, 'BACKGROUND')
 61            classes  = [ elem.replace(" ", "") for elem in classes]
 62            self.class_names = tuple(classes)
 63
 64        else:
 65            self.class_names = ('BACKGROUND',
 66            'aeroplane', 'bicycle', 'bird', 'boat',
 67            'bottle', 'bus', 'car', 'cat', 'chair',
 68            'cow', 'diningtable', 'dog', 'horse',
 69            'motorbike', 'person', 'pottedplant',
 70            'sheep', 'sofa', 'train', 'tvmonitor')
 71
 72
 73        self.class_dict = {class_name: i for i, class_name in enumerate(self.class_names)}
 74
 75    def __len__(self):
 76        return len(self.ids)
 77
 78    def __getitem__(self, index):
 79        image_id = self.ids[index]
 80        boxes, labels, is_difficult = self._get_annotation(image_id)
 81        if not self.keep_difficult:
 82            boxes = boxes[is_difficult == 0]
 83            labels = labels[is_difficult == 0]
 84        image = self._read_image(image_id)
 85        if self.transform:
 86            image, boxes, labels = self.transform(image, boxes, labels)
 87        if self.target_transform:
 88            boxes, labels = self.target_transform(boxes, labels)
 89        return image, boxes, labels
 90
 91    def get_image(self, index):
 92        image_id = self.ids[index]
 93        image = self._read_image(image_id)
 94        if self.transform:
 95            image, _ = self.transform(image)
 96        return image
 97
 98    def get_annotation(self, index):
 99        image_id = self.ids[index]
100        return image_id, self._get_annotation(image_id)
101
102    @staticmethod
103    def _read_image_ids(image_sets_file):
104        ids = []
105        with open(image_sets_file) as f:
106            for line in f:
107                ids.append(line.rstrip())
108        return ids
109
110    def _get_annotation(self, image_id):
111        annotation_file = self.root / f"Annotations/{image_id}.xml"
112        objects = ET.parse(annotation_file).findall("object")
113        boxes = []
114        labels = []
115        is_difficult = []
116        for object in objects:
117            class_name = object.find('name').text.lower().strip()
118            # we're only concerned with clases in our list
119            if class_name in self.class_dict:
120                bbox = object.find('bndbox')
121
122                # VOC dataset format follows Matlab, in which indexes start from 0
123                x1 = float(bbox.find('xmin').text) - 1
124                y1 = float(bbox.find('ymin').text) - 1
125                x2 = float(bbox.find('xmax').text) - 1
126                y2 = float(bbox.find('ymax').text) - 1
127                boxes.append([x1, y1, x2, y2])
128
129                labels.append(self.class_dict[class_name])
130                is_difficult_str = object.find('difficult').text
131                is_difficult.append(int(is_difficult_str) if is_difficult_str else 0)
132
133        return (np.array(boxes, dtype=np.float32),
134                np.array(labels, dtype=np.int64),
135                np.array(is_difficult, dtype=np.uint8))
136
137    def _read_image(self, image_id):
138        image_file = self.root / f"JPEGImages/{image_id}.jpg"
139        image = cv2.imread(str(image_file))
140        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
141        return image