From 26738d99ebb37d8481bf77aaa598f0e496dfa870 Mon Sep 17 00:00:00 2001 From: Megvii Engine Team Date: Sun, 26 Apr 2020 18:26:26 +0800 Subject: [PATCH] feat(mge/data): add objects365 dataset GitOrigin-RevId: 8eccebfc1de6a03b549f3dae1783fe17015cd723 --- .../megengine/data/dataset/vision/__init__.py | 1 + .../data/dataset/vision/objects365.py | 496 ++++++++++++++++++ 2 files changed, 497 insertions(+) create mode 100644 python_module/megengine/data/dataset/vision/objects365.py diff --git a/python_module/megengine/data/dataset/vision/__init__.py b/python_module/megengine/data/dataset/vision/__init__.py index 32e10cd9e..dd2b0fc30 100644 --- a/python_module/megengine/data/dataset/vision/__init__.py +++ b/python_module/megengine/data/dataset/vision/__init__.py @@ -13,4 +13,5 @@ from .folder import ImageFolder from .imagenet import ImageNet from .meta_vision import VisionDataset from .mnist import MNIST +from .objects365 import Objects365 from .voc import PascalVOC diff --git a/python_module/megengine/data/dataset/vision/objects365.py b/python_module/megengine/data/dataset/vision/objects365.py new file mode 100644 index 000000000..a2a50a440 --- /dev/null +++ b/python_module/megengine/data/dataset/vision/objects365.py @@ -0,0 +1,496 @@ +# -*- coding: utf-8 -*- +# MegEngine is Licensed under the Apache License, Version 2.0 (the "License") +# +# Copyright (c) 2014-2020 Megvii Inc. All rights reserved. +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# --------------------------------------------------------------------- +# Part of the following code in this file refs to maskrcnn-benchmark +# MIT License +# +# Copyright (c) 2018 Facebook +# --------------------------------------------------------------------- +import json +import os +from collections import defaultdict + +import cv2 +import numpy as np + +from .meta_vision import VisionDataset + + +class Objects365(VisionDataset): + r"""`Objects365 `_ Dataset. + """ + + supported_order = ( + "image", + "boxes", + "boxes_category", + "info", + ) + + def __init__( + self, root, ann_file, remove_images_without_annotations=False, *, order=None + ): + super().__init__(root, order=order, supported_order=self.supported_order) + + with open(ann_file, "r") as f: + dataset = json.load(f) + + self.imgs = dict() + for img in dataset["images"]: + self.imgs[img["id"]] = img + + self.img_to_anns = defaultdict(list) + for ann in dataset["annotations"]: + # for saving memory + if ( + "boxes" not in self.order + and "boxes_category" not in self.order + and "bbox" in ann + ): + del ann["bbox"] + self.img_to_anns[ann["image_id"]].append(ann) + + self.cats = dict() + for cat in dataset["categories"]: + self.cats[cat["id"]] = cat + + self.ids = list(sorted(self.imgs.keys())) + + # filter images without detection annotations + if remove_images_without_annotations: + ids = [] + for img_id in self.ids: + anno = self.img_to_anns[img_id] + # filter crowd annotations + anno = [obj for obj in anno if obj["iscrowd"] == 0] + anno = [obj for obj in anno if obj["bbox"][2] > 0 and obj["bbox"][3] > 0] + if len(anno) > 0: + ids.append(img_id) + self.img_to_anns[img_id] = anno + else: + del self.imgs[img_id] + del self.img_to_anns[img_id] + self.ids = ids + + self.json_category_id_to_contiguous_id = { + v: i + 1 for i, v in enumerate(self.cats.keys()) + } + + self.contiguous_category_id_to_json_id = { + v: k for k, v in self.json_category_id_to_contiguous_id.items() + } + + def __getitem__(self, index): + img_id = self.ids[index] + anno = self.img_to_anns[img_id] + + target = [] + for k in self.order: + if k == "image": + file_name = self.imgs[img_id]["file_name"] + path = os.path.join(self.root, file_name) + image = cv2.imread(path, cv2.IMREAD_COLOR) + target.append(image) + elif k == "boxes": + boxes = [obj["bbox"] for obj in anno] + boxes = np.array(boxes, dtype=np.float32).reshape(-1, 4) + # transfer boxes from xywh to xyxy + boxes[:, 2:] += boxes[:, :2] + target.append(boxes) + elif k == "boxes_category": + boxes_category = [obj["category_id"] for obj in anno] + boxes_category = [ + self.json_category_id_to_contiguous_id[c] for c in boxes_category + ] + boxes_category = np.array(boxes_category, dtype=np.int32) + target.append(boxes_category) + elif k == "info": + info = self.imgs[img_id] + info = [info["height"], info["width"], info["file_name"]] + target.append(info) + else: + raise NotImplementedError + + return tuple(target) + + def __len__(self): + return len(self.ids) + + def get_img_info(self, index): + img_id = self.ids[index] + img_info = self.imgs[img_id] + return img_info + + class_name = ( + "person", + "sneakers", + "chair", + "hat", + "lamp", + "bottle", + "cabinet/shelf", + "cup", + "car", + "glasses", + "picture/frame", + "desk", + "handbag", + "street lights", + "book", + "plate", + "helmet", + "leather shoes", + "pillow", + "glove", + "potted plant", + "bracelet", + "flower", + "tv", + "storage box", + "vase", + "bench", + "wine glass", + "boots", + "bowl", + "dining table", + "umbrella", + "boat", + "flag", + "speaker", + "trash bin/can", + "stool", + "backpack", + "couch", + "belt", + "carpet", + "basket", + "towel/napkin", + "slippers", + "barrel/bucket", + "coffee table", + "suv", + "toy", + "tie", + "bed", + "traffic light", + "pen/pencil", + "microphone", + "sandals", + "canned", + "necklace", + "mirror", + "faucet", + "bicycle", + "bread", + "high heels", + "ring", + "van", + "watch", + "sink", + "horse", + "fish", + "apple", + "camera", + "candle", + "teddy bear", + "cake", + "motorcycle", + "wild bird", + "laptop", + "knife", + "traffic sign", + "cell phone", + "paddle", + "truck", + "cow", + "power outlet", + "clock", + "drum", + "fork", + "bus", + "hanger", + "nightstand", + "pot/pan", + "sheep", + "guitar", + "traffic cone", + "tea pot", + "keyboard", + "tripod", + "hockey", + "fan", + "dog", + "spoon", + "blackboard/whiteboard", + "balloon", + "air conditioner", + "cymbal", + "mouse", + "telephone", + "pickup truck", + "orange", + "banana", + "airplane", + "luggage", + "skis", + "soccer", + "trolley", + "oven", + "remote", + "baseball glove", + "paper towel", + "refrigerator", + "train", + "tomato", + "machinery vehicle", + "tent", + "shampoo/shower gel", + "head phone", + "lantern", + "donut", + "cleaning products", + "sailboat", + "tangerine", + "pizza", + "kite", + "computer box", + "elephant", + "toiletries", + "gas stove", + "broccoli", + "toilet", + "stroller", + "shovel", + "baseball bat", + "microwave", + "skateboard", + "surfboard", + "surveillance camera", + "gun", + "life saver", + "cat", + "lemon", + "liquid soap", + "zebra", + "duck", + "sports car", + "giraffe", + "pumpkin", + "piano", + "stop sign", + "radiator", + "converter", + "tissue ", + "carrot", + "washing machine", + "vent", + "cookies", + "cutting/chopping board", + "tennis racket", + "candy", + "skating and skiing shoes", + "scissors", + "folder", + "baseball", + "strawberry", + "bow tie", + "pigeon", + "pepper", + "coffee machine", + "bathtub", + "snowboard", + "suitcase", + "grapes", + "ladder", + "pear", + "american football", + "basketball", + "potato", + "paint brush", + "printer", + "billiards", + "fire hydrant", + "goose", + "projector", + "sausage", + "fire extinguisher", + "extension cord", + "facial mask", + "tennis ball", + "chopsticks", + "electronic stove and gas stove", + "pie", + "frisbee", + "kettle", + "hamburger", + "golf club", + "cucumber", + "clutch", + "blender", + "tong", + "slide", + "hot dog", + "toothbrush", + "facial cleanser", + "mango", + "deer", + "egg", + "violin", + "marker", + "ship", + "chicken", + "onion", + "ice cream", + "tape", + "wheelchair", + "plum", + "bar soap", + "scale", + "watermelon", + "cabbage", + "router/modem", + "golf ball", + "pine apple", + "crane", + "fire truck", + "peach", + "cello", + "notepaper", + "tricycle", + "toaster", + "helicopter", + "green beans", + "brush", + "carriage", + "cigar", + "earphone", + "penguin", + "hurdle", + "swing", + "radio", + "CD", + "parking meter", + "swan", + "garlic", + "french fries", + "horn", + "avocado", + "saxophone", + "trumpet", + "sandwich", + "cue", + "kiwi fruit", + "bear", + "fishing rod", + "cherry", + "tablet", + "green vegetables", + "nuts", + "corn", + "key", + "screwdriver", + "globe", + "broom", + "pliers", + "volleyball", + "hammer", + "eggplant", + "trophy", + "dates", + "board eraser", + "rice", + "tape measure/ruler", + "dumbbell", + "hamimelon", + "stapler", + "camel", + "lettuce", + "goldfish", + "meat balls", + "medal", + "toothpaste", + "antelope", + "shrimp", + "rickshaw", + "trombone", + "pomegranate", + "coconut", + "jellyfish", + "mushroom", + "calculator", + "treadmill", + "butterfly", + "egg tart", + "cheese", + "pig", + "pomelo", + "race car", + "rice cooker", + "tuba", + "crosswalk sign", + "papaya", + "hair drier", + "green onion", + "chips", + "dolphin", + "sushi", + "urinal", + "donkey", + "electric drill", + "spring rolls", + "tortoise/turtle", + "parrot", + "flute", + "measuring cup", + "shark", + "steak", + "poker card", + "binoculars", + "llama", + "radish", + "noodles", + "yak", + "mop", + "crab", + "microscope", + "barbell", + "bread/bun", + "baozi", + "lion", + "red cabbage", + "polar bear", + "lighter", + "seal", + "mangosteen", + "comb", + "eraser", + "pitaya", + "scallop", + "pencil case", + "saw", + "table tennis paddle", + "okra", + "starfish", + "eagle", + "monkey", + "durian", + "game board", + "rabbit", + "french horn", + "ambulance", + "asparagus", + "hoverboard", + "pasta", + "target", + "hotair balloon", + "chainsaw", + "lobster", + "iron", + "flashlight", + ) -- GitLab