| |
| |
|
|
| import os |
| import cv2 |
| import numpy as np |
| from loguru import logger |
| from functools import wraps |
| from pycocotools.coco import COCO |
| from torch.utils.data.dataset import Dataset as torchDataset |
|
|
| COCO_CLASSES = ( |
| 'person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus', 'train', 'truck', 'boat', 'traffic light', 'fire hydrant', |
| 'stop sign', 'parking meter', 'bench', 'bird', 'cat', 'dog', 'horse', 'sheep', 'cow', 'elephant', 'bear', 'zebra', |
| 'giraffe', 'backpack', 'umbrella', 'handbag', 'tie', 'suitcase', 'frisbee', 'skis', 'snowboard', 'sports ball', 'kite', |
| 'baseball bat', 'baseball glove', 'skateboard', 'surfboard', 'tennis racket', 'bottle', 'wine glass', 'cup', 'fork', |
| 'knife', 'spoon', 'bowl', 'banana', 'apple', 'sandwich', 'orange', 'broccoli', 'carrot', 'hot dog', 'pizza', 'donut', |
| 'cake', 'chair', 'couch', 'potted plant', 'bed', 'dining table', 'toilet', 'tv', 'laptop', 'mouse', 'remote', |
| 'keyboard', 'cell phone', 'microwave', 'oven', 'toaster', 'sink', 'refrigerator', 'book', 'clock', 'vase', 'scissors', |
| 'teddy bear', 'hair drier', 'toothbrush') |
|
|
|
|
| def remove_useless_info(coco): |
| """ |
| Remove useless info in coco dataset. COCO object is modified inplace. |
| This function is mainly used for saving memory (save about 30% mem). |
| """ |
| if isinstance(coco, COCO): |
| dataset = coco.dataset |
| dataset.pop("info", None) |
| dataset.pop("licenses", None) |
| for img in dataset["images"]: |
| img.pop("license", None) |
| img.pop("coco_url", None) |
| img.pop("date_captured", None) |
| img.pop("flickr_url", None) |
| if "annotations" in coco.dataset: |
| for anno in coco.dataset["annotations"]: |
| anno.pop("segmentation", None) |
|
|
|
|
| class Dataset(torchDataset): |
| """ This class is a subclass of the base :class:`torch.utils.data.Dataset`, |
| that enables on the fly resizing of the ``input_dim``. |
| |
| Args: |
| input_dimension (tuple): (width,height) tuple with default dimensions of the network |
| """ |
|
|
| def __init__(self, input_dimension, mosaic=True): |
| super().__init__() |
| self.__input_dim = input_dimension[:2] |
| self.enable_mosaic = mosaic |
|
|
| @property |
| def input_dim(self): |
| """ |
| Dimension that can be used by transforms to set the correct image size, etc. |
| This allows transforms to have a single source of truth |
| for the input dimension of the network. |
| |
| Return: |
| list: Tuple containing the current width,height |
| """ |
| if hasattr(self, "_input_dim"): |
| return self._input_dim |
| return self.__input_dim |
|
|
| @staticmethod |
| def mosaic_getitem(getitem_fn): |
| """ |
| Decorator method that needs to be used around the ``__getitem__`` method. |br| |
| This decorator enables the closing mosaic |
| |
| Example: |
| >>> class CustomSet(ln.data.Dataset): |
| ... def __len__(self): |
| ... return 10 |
| ... @ln.data.Dataset.mosaic_getitem |
| ... def __getitem__(self, index): |
| ... return self.enable_mosaic |
| """ |
|
|
| @wraps(getitem_fn) |
| def wrapper(self, index): |
| if not isinstance(index, int): |
| self.enable_mosaic = index[0] |
| index = index[1] |
| ret_val = getitem_fn(self, index) |
| return ret_val |
|
|
| return wrapper |
|
|
|
|
| class COCODataset(Dataset): |
| """ |
| COCO dataset class. |
| """ |
|
|
| def __init__( |
| self, |
| data_dir='data/COCO', |
| json_file="instances_train2017.json", |
| name="train2017", |
| img_size=(416, 416), |
| preproc=None |
| ): |
| """ |
| COCO dataset initialization. Annotation data are read into memory by COCO API. |
| Args: |
| data_dir (str): dataset root directory |
| json_file (str): COCO json file name |
| name (str): COCO data name (e.g. 'train2017' or 'val2017') |
| img_size (tuple(int)): target image size after pre-processing |
| preproc: data augmentation strategy |
| """ |
| super().__init__(img_size) |
| self.data_dir = data_dir |
| self.json_file = json_file |
| self.coco = COCO(os.path.join(self.data_dir, "annotations", self.json_file)) |
| remove_useless_info(self.coco) |
| self.ids = self.coco.getImgIds() |
| self.class_ids = sorted(self.coco.getCatIds()) |
| self.cats = self.coco.loadCats(self.coco.getCatIds()) |
| self._classes = tuple([c["name"] for c in self.cats]) |
| self.imgs = None |
| self.name = name |
| self.img_size = img_size |
| self.preproc = preproc |
| self.annotations = self._load_coco_annotations() |
|
|
| def __len__(self): |
| return len(self.ids) |
|
|
| def __del__(self): |
| del self.imgs |
|
|
| def _load_coco_annotations(self): |
| return [self.load_anno_from_ids(_ids) for _ids in self.ids] |
|
|
| def load_anno_from_ids(self, id_): |
| im_ann = self.coco.loadImgs(id_)[0] |
| width = im_ann["width"] |
| height = im_ann["height"] |
| anno_ids = self.coco.getAnnIds(imgIds=[int(id_)], iscrowd=False) |
| annotations = self.coco.loadAnns(anno_ids) |
| objs = [] |
| for obj in annotations: |
| x1 = np.max((0, obj["bbox"][0])) |
| y1 = np.max((0, obj["bbox"][1])) |
| x2 = np.min((width, x1 + np.max((0, obj["bbox"][2])))) |
| y2 = np.min((height, y1 + np.max((0, obj["bbox"][3])))) |
| if obj["area"] > 0 and x2 >= x1 and y2 >= y1: |
| obj["clean_bbox"] = [x1, y1, x2, y2] |
| objs.append(obj) |
| num_objs = len(objs) |
| res = np.zeros((num_objs, 5)) |
| for ix, obj in enumerate(objs): |
| cls = self.class_ids.index(obj["category_id"]) |
| res[ix, 0:4] = obj["clean_bbox"] |
| res[ix, 4] = cls |
| r = min(self.img_size[0] / height, self.img_size[1] / width) |
| res[:, :4] *= r |
| img_info = (height, width) |
| resized_info = (int(height * r), int(width * r)) |
| file_name = ( |
| im_ann["file_name"] |
| if "file_name" in im_ann |
| else "{:012}".format(id_) + ".jpg" |
| ) |
| return res, img_info, resized_info, file_name |
|
|
| def load_anno(self, index): |
| return self.annotations[index][0] |
|
|
| def load_resized_img(self, index): |
| img = self.load_image(index) |
| r = min(self.img_size[0] / img.shape[0], self.img_size[1] / img.shape[1]) |
| resized_img = cv2.resize( |
| img, |
| (int(img.shape[1] * r), int(img.shape[0] * r)), |
| interpolation=cv2.INTER_LINEAR, |
| ).astype(np.uint8) |
| return resized_img |
|
|
| def load_image(self, index): |
| file_name = self.annotations[index][3] |
| img_file = os.path.join(self.data_dir, self.name, file_name) |
| img = cv2.imread(img_file) |
| assert img is not None, f"file named {img_file} not found" |
| return img |
|
|
| def pull_item(self, index): |
| id_ = self.ids[index] |
| res, img_info, resized_info, _ = self.annotations[index] |
| if self.imgs is not None: |
| pad_img = self.imgs[index] |
| img = pad_img[: resized_info[0], : resized_info[1], :].copy() |
| else: |
| img = self.load_resized_img(index) |
| return img, res.copy(), img_info, np.array([id_]) |
|
|
| @Dataset.mosaic_getitem |
| def __getitem__(self, index): |
| """ |
| One image / label pair for the given index is picked up and pre-processed. |
| |
| Args: |
| index (int): data index |
| |
| Returns: |
| img (numpy.ndarray): pre-processed image |
| target (torch.Tensor): pre-processed label data. |
| The shape is :math:`[max_labels, 5]`. |
| each label consists of [class, xc, yc, w, h]: |
| class (float): class index. |
| xc, yc (float) : center of bbox whose values range from 0 to 1. |
| w, h (float) : size of bbox whose values range from 0 to 1. |
| img_info : tuple of h, w. |
| h, w (int): original shape of the image |
| img_id (int): same as the input index. Used for evaluation. |
| """ |
| img, target, img_info, img_id = self.pull_item(index) |
| if self.preproc is not None: |
| img, target = self.preproc(img, target, self.input_dim) |
| return img, target, img_info, img_id |
|
|
