| |
| """ |
| COCO dataset which returns image_id for evaluation. |
| |
| Mostly copy-paste from https://github.com/pytorch/vision/blob/13b35ff/references/detection/coco_utils.py |
| """ |
| if __name__ == "__main__": |
| |
| import os, sys |
|
|
| sys.path.append(os.path.dirname(sys.path[0])) |
| from torchvision.datasets.vision import VisionDataset |
|
|
| import json |
| from pathlib import Path |
| import random |
| import os |
| from typing import Any, Callable, List, Optional, Tuple |
|
|
| from PIL import Image |
|
|
| import torch |
| import torch.utils.data |
| import torchvision |
| from pycocotools import mask as coco_mask |
|
|
| from datasets.data_util import preparing_dataset |
| import datasets.transforms as T |
| from util.box_ops import box_cxcywh_to_xyxy, box_iou |
|
|
| __all__ = ["build"] |
|
|
|
|
| class label2compat: |
| def __init__(self) -> None: |
| self.category_map_str = { |
| "1": 1, |
| "2": 2, |
| "3": 3, |
| "4": 4, |
| "5": 5, |
| "6": 6, |
| "7": 7, |
| "8": 8, |
| "9": 9, |
| "10": 10, |
| "11": 11, |
| "13": 12, |
| "14": 13, |
| "15": 14, |
| "16": 15, |
| "17": 16, |
| "18": 17, |
| "19": 18, |
| "20": 19, |
| "21": 20, |
| "22": 21, |
| "23": 22, |
| "24": 23, |
| "25": 24, |
| "27": 25, |
| "28": 26, |
| "31": 27, |
| "32": 28, |
| "33": 29, |
| "34": 30, |
| "35": 31, |
| "36": 32, |
| "37": 33, |
| "38": 34, |
| "39": 35, |
| "40": 36, |
| "41": 37, |
| "42": 38, |
| "43": 39, |
| "44": 40, |
| "46": 41, |
| "47": 42, |
| "48": 43, |
| "49": 44, |
| "50": 45, |
| "51": 46, |
| "52": 47, |
| "53": 48, |
| "54": 49, |
| "55": 50, |
| "56": 51, |
| "57": 52, |
| "58": 53, |
| "59": 54, |
| "60": 55, |
| "61": 56, |
| "62": 57, |
| "63": 58, |
| "64": 59, |
| "65": 60, |
| "67": 61, |
| "70": 62, |
| "72": 63, |
| "73": 64, |
| "74": 65, |
| "75": 66, |
| "76": 67, |
| "77": 68, |
| "78": 69, |
| "79": 70, |
| "80": 71, |
| "81": 72, |
| "82": 73, |
| "84": 74, |
| "85": 75, |
| "86": 76, |
| "87": 77, |
| "88": 78, |
| "89": 79, |
| "90": 80, |
| } |
| self.category_map = {int(k): v for k, v in self.category_map_str.items()} |
|
|
| def __call__(self, target, img=None): |
| labels = target["labels"] |
| res = torch.zeros(labels.shape, dtype=labels.dtype) |
| for idx, item in enumerate(labels): |
| res[idx] = self.category_map[item.item()] - 1 |
| target["label_compat"] = res |
| if img is not None: |
| return target, img |
| else: |
| return target |
|
|
|
|
| class label_compat2onehot: |
| def __init__(self, num_class=80, num_output_objs=1): |
| self.num_class = num_class |
| self.num_output_objs = num_output_objs |
| if num_output_objs != 1: |
| raise DeprecationWarning( |
| "num_output_objs!=1, which is only used for comparison" |
| ) |
|
|
| def __call__(self, target, img=None): |
| labels = target["label_compat"] |
| place_dict = {k: 0 for k in range(self.num_class)} |
| if self.num_output_objs == 1: |
| res = torch.zeros(self.num_class) |
| for i in labels: |
| itm = i.item() |
| res[itm] = 1.0 |
| else: |
| |
| res = torch.zeros(self.num_class, self.num_output_objs) |
| for i in labels: |
| itm = i.item() |
| res[itm][place_dict[itm]] = 1.0 |
| place_dict[itm] += 1 |
| target["label_compat_onehot"] = res |
| if img is not None: |
| return target, img |
| else: |
| return target |
|
|
|
|
| class box_label_catter: |
| def __init__(self): |
| pass |
|
|
| def __call__(self, target, img=None): |
| labels = target["label_compat"] |
| boxes = target["boxes"] |
| box_label = torch.cat((boxes, labels.unsqueeze(-1)), 1) |
| target["box_label"] = box_label |
| if img is not None: |
| return target, img |
| else: |
| return target |
|
|
|
|
| class RandomSelectBoxlabels: |
| def __init__( |
| self, |
| num_classes, |
| leave_one_out=False, |
| blank_prob=0.8, |
| prob_first_item=0.0, |
| prob_random_item=0.0, |
| prob_last_item=0.8, |
| prob_stop_sign=0.2, |
| ) -> None: |
| self.num_classes = num_classes |
| self.leave_one_out = leave_one_out |
| self.blank_prob = blank_prob |
|
|
| self.set_state( |
| prob_first_item, prob_random_item, prob_last_item, prob_stop_sign |
| ) |
|
|
| def get_state(self): |
| return [ |
| self.prob_first_item, |
| self.prob_random_item, |
| self.prob_last_item, |
| self.prob_stop_sign, |
| ] |
|
|
| def set_state( |
| self, prob_first_item, prob_random_item, prob_last_item, prob_stop_sign |
| ): |
| sum_prob = prob_first_item + prob_random_item + prob_last_item + prob_stop_sign |
| assert sum_prob - 1 < 1e-6, ( |
| f"Sum up all prob = {sum_prob}. prob_first_item:{prob_first_item}" |
| + f"prob_random_item:{prob_random_item}, prob_last_item:{prob_last_item}" |
| + f"prob_stop_sign:{prob_stop_sign}" |
| ) |
|
|
| self.prob_first_item = prob_first_item |
| self.prob_random_item = prob_random_item |
| self.prob_last_item = prob_last_item |
| self.prob_stop_sign = prob_stop_sign |
|
|
| def sample_for_pred_first_item(self, box_label: torch.FloatTensor): |
| box_label_known = torch.Tensor(0, 5) |
| box_label_unknown = box_label |
| return box_label_known, box_label_unknown |
|
|
| def sample_for_pred_random_item(self, box_label: torch.FloatTensor): |
| n_select = int(random.random() * box_label.shape[0]) |
| box_label = box_label[torch.randperm(box_label.shape[0])] |
| box_label_known = box_label[:n_select] |
| box_label_unknown = box_label[n_select:] |
| return box_label_known, box_label_unknown |
|
|
| def sample_for_pred_last_item(self, box_label: torch.FloatTensor): |
| box_label_perm = box_label[torch.randperm(box_label.shape[0])] |
| known_label_list = [] |
| box_label_known = [] |
| box_label_unknown = [] |
| for item in box_label_perm: |
| label_i = item[4].item() |
| if label_i in known_label_list: |
| box_label_known.append(item) |
| else: |
| |
| box_label_unknown.append(item) |
| known_label_list.append(label_i) |
| box_label_known = ( |
| torch.stack(box_label_known) |
| if len(box_label_known) > 0 |
| else torch.Tensor(0, 5) |
| ) |
| box_label_unknown = ( |
| torch.stack(box_label_unknown) |
| if len(box_label_unknown) > 0 |
| else torch.Tensor(0, 5) |
| ) |
| return box_label_known, box_label_unknown |
|
|
| def sample_for_pred_stop_sign(self, box_label: torch.FloatTensor): |
| box_label_unknown = torch.Tensor(0, 5) |
| box_label_known = box_label |
| return box_label_known, box_label_unknown |
|
|
| def __call__(self, target, img=None): |
| box_label = target["box_label"] |
|
|
| dice_number = random.random() |
|
|
| if dice_number < self.prob_first_item: |
| box_label_known, box_label_unknown = self.sample_for_pred_first_item( |
| box_label |
| ) |
| elif dice_number < self.prob_first_item + self.prob_random_item: |
| box_label_known, box_label_unknown = self.sample_for_pred_random_item( |
| box_label |
| ) |
| elif ( |
| dice_number |
| < self.prob_first_item + self.prob_random_item + self.prob_last_item |
| ): |
| box_label_known, box_label_unknown = self.sample_for_pred_last_item( |
| box_label |
| ) |
| else: |
| box_label_known, box_label_unknown = self.sample_for_pred_stop_sign( |
| box_label |
| ) |
|
|
| target["label_onehot_known"] = label2onehot( |
| box_label_known[:, -1], self.num_classes |
| ) |
| target["label_onehot_unknown"] = label2onehot( |
| box_label_unknown[:, -1], self.num_classes |
| ) |
| target["box_label_known"] = box_label_known |
| target["box_label_unknown"] = box_label_unknown |
|
|
| return target, img |
|
|
|
|
| class RandomDrop: |
| def __init__(self, p=0.2) -> None: |
| self.p = p |
|
|
| def __call__(self, target, img=None): |
| known_box = target["box_label_known"] |
| num_known_box = known_box.size(0) |
| idxs = torch.rand(num_known_box) |
| |
| target["box_label_known"] = known_box[idxs > self.p] |
| return target, img |
|
|
|
|
| class BboxPertuber: |
| def __init__(self, max_ratio=0.02, generate_samples=1000) -> None: |
| self.max_ratio = max_ratio |
| self.generate_samples = generate_samples |
| self.samples = self.generate_pertube_samples() |
| self.idx = 0 |
|
|
| def generate_pertube_samples(self): |
| import torch |
|
|
| samples = (torch.rand(self.generate_samples, 5) - 0.5) * 2 * self.max_ratio |
| return samples |
|
|
| def __call__(self, target, img): |
| known_box = target["box_label_known"] |
| K = known_box.shape[0] |
| known_box_pertube = torch.zeros(K, 6) |
| if K == 0: |
| pass |
| else: |
| if self.idx + K > self.generate_samples: |
| self.idx = 0 |
| delta = self.samples[self.idx : self.idx + K, :] |
| known_box_pertube[:, :4] = known_box[:, :4] + delta[:, :4] |
| iou = ( |
| torch.diag( |
| box_iou( |
| box_cxcywh_to_xyxy(known_box[:, :4]), |
| box_cxcywh_to_xyxy(known_box_pertube[:, :4]), |
| )[0] |
| ) |
| ) * (1 + delta[:, -1]) |
| known_box_pertube[:, 4].copy_(iou) |
| known_box_pertube[:, -1].copy_(known_box[:, -1]) |
|
|
| target["box_label_known_pertube"] = known_box_pertube |
| return target, img |
|
|
|
|
| class RandomCutout: |
| def __init__(self, factor=0.5) -> None: |
| self.factor = factor |
|
|
| def __call__(self, target, img=None): |
| unknown_box = target["box_label_unknown"] |
| known_box = target["box_label_known_pertube"] |
| Ku = unknown_box.size(0) |
|
|
| known_box_add = torch.zeros(Ku, 6) |
| known_box_add[:, :5] = unknown_box |
| known_box_add[:, 5].uniform_(0.5, 1) |
|
|
| known_box_add[:, :2] += known_box_add[:, 2:4] * (torch.rand(Ku, 2) - 0.5) / 2 |
| known_box_add[:, 2:4] /= 2 |
|
|
| target["box_label_known_pertube"] = torch.cat((known_box, known_box_add)) |
| return target, img |
|
|
|
|
| class RandomSelectBoxes: |
| def __init__(self, num_class=80) -> None: |
| Warning("This is such a slow function and will be deprecated soon!!!") |
| self.num_class = num_class |
|
|
| def __call__(self, target, img=None): |
| boxes = target["boxes"] |
| labels = target["label_compat"] |
|
|
| |
| boxs_list = [[] for i in range(self.num_class)] |
| for idx, item in enumerate(boxes): |
| label = labels[idx].item() |
| boxs_list[label].append(item) |
| boxs_list_tensor = [ |
| torch.stack(i) if len(i) > 0 else torch.Tensor(0, 4) for i in boxs_list |
| ] |
|
|
| |
| box_known = [] |
| box_unknown = [] |
| for idx, item in enumerate(boxs_list_tensor): |
| ncnt = item.shape[0] |
| nselect = int( |
| random.random() * ncnt |
| ) |
|
|
| item = item[torch.randperm(ncnt)] |
| |
| box_known.append(item[:nselect]) |
| box_unknown.append(item[nselect:]) |
|
|
| |
| |
| |
| target["known_box"] = box_known |
| target["unknown_box"] = box_unknown |
| return target, img |
|
|
|
|
| def label2onehot(label, num_classes): |
| """ |
| label: Tensor(K) |
| """ |
| res = torch.zeros(num_classes) |
| for i in label: |
| itm = int(i.item()) |
| res[itm] = 1.0 |
| return res |
|
|
|
|
| class MaskCrop: |
| def __init__(self) -> None: |
| pass |
|
|
| def __call__(self, target, img): |
| known_box = target["known_box"] |
| h, w = img.shape[1:] |
| |
|
|
| scale = torch.Tensor([w, h, w, h]) |
|
|
| |
| for boxes in known_box: |
| if boxes.shape[0] == 0: |
| continue |
| box_xyxy = box_cxcywh_to_xyxy(boxes) * scale |
| for box in box_xyxy: |
| x1, y1, x2, y2 = [int(i) for i in box.tolist()] |
| img[:, y1:y2, x1:x2] = 0 |
| |
| |
| return target, img |
|
|
|
|
| dataset_hook_register = { |
| "label2compat": label2compat, |
| "label_compat2onehot": label_compat2onehot, |
| "box_label_catter": box_label_catter, |
| "RandomSelectBoxlabels": RandomSelectBoxlabels, |
| "RandomSelectBoxes": RandomSelectBoxes, |
| "MaskCrop": MaskCrop, |
| "BboxPertuber": BboxPertuber, |
| } |
|
|
|
|
| class CocoDetection(torchvision.datasets.CocoDetection): |
| def __init__( |
| self, img_folder, ann_file, transforms, return_masks, aux_target_hacks=None |
| ): |
| super(CocoDetection, self).__init__(img_folder, ann_file) |
| self._transforms = transforms |
| self.prepare = ConvertCocoPolysToMask(return_masks) |
| self.aux_target_hacks = aux_target_hacks |
|
|
| def change_hack_attr(self, hackclassname, attrkv_dict): |
| target_class = dataset_hook_register[hackclassname] |
| for item in self.aux_target_hacks: |
| if isinstance(item, target_class): |
| for k, v in attrkv_dict.items(): |
| setattr(item, k, v) |
|
|
| def get_hack(self, hackclassname): |
| target_class = dataset_hook_register[hackclassname] |
| for item in self.aux_target_hacks: |
| if isinstance(item, target_class): |
| return item |
|
|
| def _load_image(self, id: int) -> Image.Image: |
| path = self.coco.loadImgs(id)[0]["file_name"] |
| abs_path = os.path.join(self.root, path) |
| return Image.open(abs_path).convert("RGB") |
|
|
| def __getitem__(self, idx): |
| """ |
| Output: |
| - target: dict of multiple items |
| - boxes: Tensor[num_box, 4]. \ |
| Init type: x0,y0,x1,y1. unnormalized data. |
| Final type: cx,cy,w,h. normalized data. |
| """ |
| try: |
| img, target = super(CocoDetection, self).__getitem__(idx) |
|
|
| except: |
| print("Error idx: {}".format(idx)) |
| idx += 1 |
| img, target = super(CocoDetection, self).__getitem__(idx) |
|
|
| image_id = self.ids[idx] |
| target = {"image_id": image_id, "annotations": target} |
| exemp_count = 0 |
| for instance in target["annotations"]: |
| if instance["area"] != 4: |
| exemp_count += 1 |
| |
| assert exemp_count == 3 |
| img, target = self.prepare(img, target) |
| target["exemplars"] = target["boxes"][-3:] |
| |
| if image_id == 6003: |
| target["exemplars"] = torch.tensor([]) |
| target["boxes"] = target["boxes"][:-3] |
| target["labels"] = target["labels"][:-3] |
| target["labels_uncropped"] = torch.clone(target["labels"]) |
|
|
| if self._transforms is not None: |
| img, target = self._transforms(img, target) |
|
|
| |
| if self.aux_target_hacks is not None: |
| for hack_runner in self.aux_target_hacks: |
| target, img = hack_runner(target, img=img) |
|
|
| return img, target |
|
|
|
|
| def convert_coco_poly_to_mask(segmentations, height, width): |
| masks = [] |
| for polygons in segmentations: |
| rles = coco_mask.frPyObjects(polygons, height, width) |
| mask = coco_mask.decode(rles) |
| if len(mask.shape) < 3: |
| mask = mask[..., None] |
| mask = torch.as_tensor(mask, dtype=torch.uint8) |
| mask = mask.any(dim=2) |
| masks.append(mask) |
| if masks: |
| masks = torch.stack(masks, dim=0) |
| else: |
| masks = torch.zeros((0, height, width), dtype=torch.uint8) |
| return masks |
|
|
|
|
| class ConvertCocoPolysToMask(object): |
| def __init__(self, return_masks=False): |
| self.return_masks = return_masks |
|
|
| def __call__(self, image, target): |
| w, h = image.size |
|
|
| image_id = target["image_id"] |
| image_id = torch.tensor([image_id]) |
|
|
| anno = target["annotations"] |
|
|
| anno = [obj for obj in anno if "iscrowd" not in obj or obj["iscrowd"] == 0] |
|
|
| boxes = [obj["bbox"] for obj in anno] |
| |
| boxes = torch.as_tensor(boxes, dtype=torch.float32).reshape(-1, 4) |
| boxes[:, 2:] += boxes[:, :2] |
| boxes[:, 0::2].clamp_(min=0, max=w) |
| boxes[:, 1::2].clamp_(min=0, max=h) |
|
|
| classes = [obj["category_id"] for obj in anno] |
| classes = torch.tensor(classes, dtype=torch.int64) |
|
|
| if self.return_masks: |
| segmentations = [obj["segmentation"] for obj in anno] |
| masks = convert_coco_poly_to_mask(segmentations, h, w) |
|
|
| keypoints = None |
| if anno and "keypoints" in anno[0]: |
| keypoints = [obj["keypoints"] for obj in anno] |
| keypoints = torch.as_tensor(keypoints, dtype=torch.float32) |
| num_keypoints = keypoints.shape[0] |
| if num_keypoints: |
| keypoints = keypoints.view(num_keypoints, -1, 3) |
|
|
| keep = (boxes[:, 3] > boxes[:, 1]) & (boxes[:, 2] > boxes[:, 0]) |
| boxes = boxes[keep] |
| classes = classes[keep] |
| if self.return_masks: |
| masks = masks[keep] |
| if keypoints is not None: |
| keypoints = keypoints[keep] |
|
|
| target = {} |
| target["boxes"] = boxes |
| target["labels"] = classes |
| if self.return_masks: |
| target["masks"] = masks |
| target["image_id"] = image_id |
| if keypoints is not None: |
| target["keypoints"] = keypoints |
|
|
| |
| area = torch.tensor([obj["area"] for obj in anno]) |
| iscrowd = torch.tensor( |
| [obj["iscrowd"] if "iscrowd" in obj else 0 for obj in anno] |
| ) |
| target["area"] = area[keep] |
| target["iscrowd"] = iscrowd[keep] |
|
|
| target["orig_size"] = torch.as_tensor([int(h), int(w)]) |
| target["size"] = torch.as_tensor([int(h), int(w)]) |
|
|
| return image, target |
|
|
|
|
| def make_coco_transforms(image_set, fix_size=False, strong_aug=False, args=None): |
| normalize = T.Compose( |
| [T.ToTensor(), T.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])] |
| ) |
|
|
| |
| scales = [480, 512, 544, 576, 608, 640, 672, 704, 736, 768, 800] |
| max_size = 1333 |
| scales2_resize = [400, 500, 600] |
| scales2_crop = [384, 600] |
|
|
| |
| scales = getattr(args, "data_aug_scales", scales) |
| max_size = getattr(args, "data_aug_max_size", max_size) |
| scales2_resize = getattr(args, "data_aug_scales2_resize", scales2_resize) |
| scales2_crop = getattr(args, "data_aug_scales2_crop", scales2_crop) |
|
|
| |
| data_aug_scale_overlap = getattr(args, "data_aug_scale_overlap", None) |
| if data_aug_scale_overlap is not None and data_aug_scale_overlap > 0: |
| data_aug_scale_overlap = float(data_aug_scale_overlap) |
| scales = [int(i * data_aug_scale_overlap) for i in scales] |
| max_size = int(max_size * data_aug_scale_overlap) |
| scales2_resize = [int(i * data_aug_scale_overlap) for i in scales2_resize] |
| scales2_crop = [int(i * data_aug_scale_overlap) for i in scales2_crop] |
|
|
| datadict_for_print = { |
| "scales": scales, |
| "max_size": max_size, |
| "scales2_resize": scales2_resize, |
| "scales2_crop": scales2_crop, |
| } |
| |
|
|
| if image_set == "train": |
| if fix_size: |
| return T.Compose( |
| [ |
| T.RandomHorizontalFlip(), |
| T.RandomResize([(max_size, max(scales))]), |
| |
| normalize, |
| ] |
| ) |
|
|
| if strong_aug: |
| import datasets.sltransform as SLT |
|
|
| return T.Compose( |
| [ |
| T.RandomHorizontalFlip(), |
| T.RandomSelect( |
| T.RandomResize(scales, max_size=max_size), |
| T.Compose( |
| [ |
| T.RandomResize(scales2_resize), |
| T.RandomSizeCrop(*scales2_crop), |
| T.RandomResize(scales, max_size=max_size), |
| ] |
| ), |
| ), |
| SLT.RandomSelectMulti( |
| [ |
| SLT.RandomCrop(), |
| SLT.LightingNoise(), |
| SLT.AdjustBrightness(2), |
| SLT.AdjustContrast(2), |
| ] |
| ), |
| normalize, |
| ] |
| ) |
|
|
| return T.Compose( |
| [ |
| T.RandomHorizontalFlip(), |
| T.RandomSelect( |
| T.RandomResize(scales, max_size=max_size), |
| T.Compose( |
| [ |
| T.RandomResize(scales2_resize), |
| T.RandomSizeCrop(*scales2_crop), |
| T.RandomResize(scales, max_size=max_size), |
| ] |
| ), |
| ), |
| normalize, |
| ] |
| ) |
|
|
| if image_set in ["val", "eval_debug", "train_reg", "test"]: |
| if os.environ.get("GFLOPS_DEBUG_SHILONG", False) == "INFO": |
| print("Under debug mode for flops calculation only!!!!!!!!!!!!!!!!") |
| return T.Compose( |
| [ |
| T.ResizeDebug((1280, 800)), |
| normalize, |
| ] |
| ) |
|
|
| print("max(scales): " + str(max(scales))) |
|
|
| return T.Compose( |
| [ |
| T.RandomResize([max(scales)], max_size=max_size), |
| normalize, |
| ] |
| ) |
|
|
| raise ValueError(f"unknown {image_set}") |
|
|
|
|
| def get_aux_target_hacks_list(image_set, args): |
| if args.modelname in ["q2bs_mask", "q2bs"]: |
| aux_target_hacks_list = [ |
| label2compat(), |
| label_compat2onehot(), |
| RandomSelectBoxes(num_class=args.num_classes), |
| ] |
| if args.masked_data and image_set == "train": |
| |
| aux_target_hacks_list.append(MaskCrop()) |
| elif args.modelname in [ |
| "q2bm_v2", |
| "q2bs_ce", |
| "q2op", |
| "q2ofocal", |
| "q2opclip", |
| "q2ocqonly", |
| ]: |
| aux_target_hacks_list = [ |
| label2compat(), |
| label_compat2onehot(), |
| box_label_catter(), |
| RandomSelectBoxlabels( |
| num_classes=args.num_classes, |
| prob_first_item=args.prob_first_item, |
| prob_random_item=args.prob_random_item, |
| prob_last_item=args.prob_last_item, |
| prob_stop_sign=args.prob_stop_sign, |
| ), |
| BboxPertuber(max_ratio=0.02, generate_samples=1000), |
| ] |
| elif args.modelname in ["q2omask", "q2osa"]: |
| if args.coco_aug: |
| aux_target_hacks_list = [ |
| label2compat(), |
| label_compat2onehot(), |
| box_label_catter(), |
| RandomSelectBoxlabels( |
| num_classes=args.num_classes, |
| prob_first_item=args.prob_first_item, |
| prob_random_item=args.prob_random_item, |
| prob_last_item=args.prob_last_item, |
| prob_stop_sign=args.prob_stop_sign, |
| ), |
| RandomDrop(p=0.2), |
| BboxPertuber(max_ratio=0.02, generate_samples=1000), |
| RandomCutout(factor=0.5), |
| ] |
| else: |
| aux_target_hacks_list = [ |
| label2compat(), |
| label_compat2onehot(), |
| box_label_catter(), |
| RandomSelectBoxlabels( |
| num_classes=args.num_classes, |
| prob_first_item=args.prob_first_item, |
| prob_random_item=args.prob_random_item, |
| prob_last_item=args.prob_last_item, |
| prob_stop_sign=args.prob_stop_sign, |
| ), |
| BboxPertuber(max_ratio=0.02, generate_samples=1000), |
| ] |
| else: |
| aux_target_hacks_list = None |
|
|
| return aux_target_hacks_list |
|
|
|
|
| def build(image_set, args, datasetinfo): |
| img_folder = datasetinfo["root"] |
| ann_file = datasetinfo["anno"] |
|
|
| |
| if os.environ.get("DATA_COPY_SHILONG") == "INFO": |
| preparing_dataset( |
| dict(img_folder=img_folder, ann_file=ann_file), image_set, args |
| ) |
|
|
| try: |
| strong_aug = args.strong_aug |
| except: |
| strong_aug = False |
| print(img_folder, ann_file) |
| dataset = CocoDetection( |
| img_folder, |
| ann_file, |
| transforms=make_coco_transforms( |
| image_set, fix_size=args.fix_size, strong_aug=strong_aug, args=args |
| ), |
| return_masks=args.masks, |
| aux_target_hacks=None, |
| ) |
| return dataset |
|
|
|
|
| if __name__ == "__main__": |
| |
| dataset_o365 = CocoDetection( |
| "/path/Objects365/train/", |
| "/path/Objects365/slannos/anno_preprocess_train_v2.json", |
| transforms=None, |
| return_masks=False, |
| ) |
| print("len(dataset_o365):", len(dataset_o365)) |
|
|