seg/datasets/ade_ov.py

import copy
from typing import List

from mmdet.registry import DATASETS
from mmdet.datasets.coco_panoptic import CocoPanopticDataset
from mmengine import get_local_path


@DATASETS.register_module()
class ADEPanopticOVDataset(CocoPanopticDataset):
    """ADE Open Vocabulary dataset for Panoptic segmentation.
    The class names are changed.
    """

    METAINFO = {
        'classes':
            (
                'bed,beds',
                'windowpane,window,windows',
                'cabinet,cabinets,wall mounted cabine',
                'person,child,girl,boy,woman,man,people,children,girls,boys,women,men',
                'door,double door,doors',
                'table,tables,tablecloth',
                'curtain,drape,drapery,mantle,pall',
                'chair,chairs',
                'car,automobile,cars',
                'painting,picture,paintings,pictures,wallart,framed canvas',
                'sofa,couch,sofas,couches',
                'shelf,shelves',
                'mirror,mirrors',
                'armchair,armchairs',
                'seat,seats',
                'fence,fencing',
                'desk,desks',
                'wardrobe,closet,press,wardrobes,closets',
                'lamp,lamps',
                'bathtub,bathing tub,bath,tub',
                'railing,rail',
                'cushion,cushions',
                'box,boxes',
                'column,pillar',
                'signboard,sign,signboards,signs',
                'chest of drawers,chest,bureau,dresser',
                'counter',
                'sink',
                'fireplace,hearth,open fireplace',
                'refrigerator,icebox',
                'stairs,steps',
                'case,display case,showcase,vitrine',
                'pool table,billiard table,snooker table',
                'pillow,pillows',
                'screen door,shower door',
                'bookcase',
                'coffee table,cocktail table',
                'toilet,commode,crapper,potty',
                'flower,flowers',
                'book,books',
                'bench,benches',
                'countertop,counter top,worktop',
                'stove,kitchen stove,kitchen range,kitchen range,cooking stove',
                'palm tree,palm trees',
                'kitchen island',
                'computer,computing machine,computing device,data processor,electronic computer,information processing system',
                'swivel chair',
                'boat',
                'arcade machine,arcade machines',
                'bus,autobus,double-decker,jitney,motorbus,motorcoach,omnibus,passenger vehicle',
                'towel',
                'light bulb,lightbulb,bulb,incandescent lamp,electric light,electric-light bulb',
                'truck,motortruck',
                'chandelier,pendant,pendent',
                'awning,sunshade,sunblind',
                'streetlight,street lamp',
                'booth,cubicle,stall,kiosk',
                'television receiver,television,television set,tv,tv set',
                'airplane,aeroplane,airplanes,aeroplanes',
                'apparel,wearing apparel,dress,clothes',
                'pole',
                'bannister,banister,balustrade,balusters,handrail',
                'ottoman,pouf,pouffe,puff,hassock',
                'bottle,bottles,water bottle',
                'van',
                'ship',
                'fountain',
                'washer,automatic washer,washing machine',
                'plaything,toy,toys',
                'stool,stools',
                'barrel,cask,barrels,casks',
                'basket,handbasket',
                'bag,bags,gift bag,paper bag',
                'minibike,motorbike',
                'oven',
                'ball,balls',
                'food,solid food',
                'step,stair',
                'trade name,brand name,brand,marque',
                'microwave,microwave oven',
                'plant pots,plant pot,flower pot,flowerpot,planter',
                'animal,animate being,dog,cat,horse,cow,sheep,zebra,girraffe,bird',
                'bicycle,bike',
                'dishwasher,dish washer,dishwashing machine',
                'projection screen',
                'sculpture,sculptures',
                'exhaust hood',
                'sconce,sconce lamp,sconce light',
                'vase,vases',
                'traffic light,traffic signal,traffic lights',
                'tray,trays',
                'ashcan,trash can,garbage can,wastebin,ash bin,ash-bin,ashbin,dustbin,trash barrel,trash bin',
                'ceiling fan,floor fan',
                'plate,plates',
                'monitor,monitoring device,monitors',
                'bulletin board,notice board',
                'radiator',
                'cup,cups,drinking glass,drinking glasses',
                'clock',
                'flag,flags',

                'wall,walls,brick wall,stone wall,interior wall',
                'building,buildings,edifice,edifices',
                'sky,clouds',
                'floor,flooring',
                'tree,trees',
                'ceiling',
                'road,route,street,roads,streets,routes',
                'grass,grass field',
                'sidewalk,pavement',
                'earth,ground',
                'mountain,mount,mountains',
                'plant,flora,plant life,plants,bushes',
                'water',
                'house exterior',
                'sea,ocean',
                'rug,carpet,carpeting',
                'field',
                'rock,stone,rocks,stones',
                'pedestal',
                'sand',
                'skyscraper,skyscrapers',
                'grandstand,covered stand',
                'path',
                'runway',
                'stairway,staircase',
                'river',
                'bridge,span',
                'window screen,door screen',
                'hill',
                'bar',
                'hovel,hut,hutch,shack,shanty',
                'tower,towers',
                'dirt track',
                'land,soil',
                'escalator,moving staircase,moving stairway',
                'buffet,sideboard',
                'poster,posting,placard,notice,bill,card',
                'stage',
                'conveyer belt,conveyor belt,conveyer,conveyor,transporter',
                'canopy',
                'swimming pool,swimming bath',
                'waterfall,falls',
                'tent,collapsible shelter',
                'cradle',
                'tank,storage tank',
                'lake',
                'blanket,cover',
                'pier,wharf,wharfage,dock',
                'crt screen',
                'shower',
            ),
        'thing_classes':
            (
                'bed,beds',
                'windowpane,window,windows',
                'cabinet,cabinets,wall mounted cabine',
                'person,child,girl,boy,woman,man,people,children,girls,boys,women,men',
                'door,double door,doors',
                'table,tables,tablecloth',
                'curtain,drape,drapery,mantle,pall',
                'chair,chairs',
                'car,automobile,cars',
                'painting,picture,paintings,pictures,wallart,framed canvas',
                'sofa,couch,sofas,couches',
                'shelf,shelves',
                'mirror,mirrors',
                'armchair,armchairs',
                'seat,seats',
                'fence,fencing',
                'desk,desks',
                'wardrobe,closet,press,wardrobes,closets',
                'lamp,lamps',
                'bathtub,bathing tub,bath,tub',
                'railing,rail',
                'cushion,cushions',
                'box,boxes',
                'column,pillar',
                'signboard,sign,signboards,signs',
                'chest of drawers,chest,bureau,dresser',
                'counter',
                'sink',
                'fireplace,hearth,open fireplace',
                'refrigerator,icebox',
                'stairs,steps',
                'case,display case,showcase,vitrine',
                'pool table,billiard table,snooker table',
                'pillow,pillows',
                'screen door,shower door',
                'bookcase',
                'coffee table,cocktail table',
                'toilet,commode,crapper,potty',
                'flower,flowers',
                'book,books',
                'bench,benches',
                'countertop,counter top,worktop',
                'stove,kitchen stove,kitchen range,kitchen range,cooking stove',
                'palm tree,palm trees',
                'kitchen island',
                'computer,computing machine,computing device,data processor,electronic computer,information processing system',
                'swivel chair',
                'boat',
                'arcade machine,arcade machines',
                'bus,autobus,double-decker,jitney,motorbus,motorcoach,omnibus,passenger vehicle',
                'towel',
                'light bulb,lightbulb,bulb,incandescent lamp,electric light,electric-light bulb',
                'truck,motortruck',
                'chandelier,pendant,pendent',
                'awning,sunshade,sunblind',
                'streetlight,street lamp',
                'booth,cubicle,stall,kiosk',
                'television receiver,television,television set,tv,tv set',
                'airplane,aeroplane,airplanes,aeroplanes',
                'apparel,wearing apparel,dress,clothes',
                'pole',
                'bannister,banister,balustrade,balusters,handrail',
                'ottoman,pouf,pouffe,puff,hassock',
                'bottle,bottles,water bottle',
                'van',
                'ship',
                'fountain',
                'washer,automatic washer,washing machine',
                'plaything,toy,toys',
                'stool,stools',
                'barrel,cask,barrels,casks',
                'basket,handbasket',
                'bag,bags,gift bag,paper bag',
                'minibike,motorbike',
                'oven',
                'ball,balls',
                'food,solid food',
                'step,stair',
                'trade name,brand name,brand,marque',
                'microwave,microwave oven',
                'plant pots,plant pot,flower pot,flowerpot,planter',
                'animal,animate being,dog,cat,horse,cow,sheep,zebra,girraffe,bird',
                'bicycle,bike',
                'dishwasher,dish washer,dishwashing machine',
                'projection screen',
                'sculpture,sculptures',
                'exhaust hood',
                'sconce,sconce lamp,sconce light',
                'vase,vases',
                'traffic light,traffic signal,traffic lights',
                'tray,trays',
                'ashcan,trash can,garbage can,wastebin,ash bin,ash-bin,ashbin,dustbin,trash barrel,trash bin',
                'ceiling fan,floor fan',
                'plate,plates',
                'monitor,monitoring device,monitors',
                'bulletin board,notice board',
                'radiator',
                'cup,cups,drinking glass,drinking glasses',
                'clock',
                'flag,flags',
            ),
        'stuff_classes':
            (
                'wall,walls,brick wall,stone wall,interior wall',
                'building,buildings,edifice,edifices',
                'sky,clouds',
                'floor,flooring',
                'tree,trees',
                'ceiling',
                'road,route,street,roads,streets,routes',
                'grass,grass field',
                'sidewalk,pavement',
                'earth,ground',
                'mountain,mount,mountains',
                'plant,flora,plant life,plants,bushes',
                'water',
                'house exterior',
                'sea,ocean',
                'rug,carpet,carpeting',
                'field',
                'rock,stone,rocks,stones',
                'pedestal',
                'sand',
                'skyscraper,skyscrapers',
                'grandstand,covered stand',
                'path',
                'runway',
                'stairway,staircase',
                'river',
                'bridge,span',
                'window screen,door screen',
                'hill',
                'bar',
                'hovel,hut,hutch,shack,shanty',
                'tower,towers',
                'dirt track',
                'land,soil',
                'escalator,moving staircase,moving stairway',
                'buffet,sideboard',
                'poster,posting,placard,notice,bill,card',
                'stage',
                'conveyer belt,conveyor belt,conveyer,conveyor,transporter',
                'canopy',
                'swimming pool,swimming bath',
                'waterfall,falls',
                'tent,collapsible shelter',
                'cradle',
                'tank,storage tank',
                'lake',
                'blanket,cover',
                'pier,wharf,wharfage,dock',
                'crt screen',
                'shower',
            ),
    }

    def load_data_list(self) -> List[dict]:
        """Load annotations from an annotation file named as ``self.ann_file``

        Returns:
            List[dict]: A list of annotation.
        """  # noqa: E501
        with get_local_path(
                self.ann_file, backend_args=self.backend_args) as local_path:
            self.coco = self.COCOAPI(local_path)
            for idx, name in enumerate(self.metainfo['classes']):
                if not (self.coco.cats[idx]['name'].strip() in name.split(',')):
                    print(f"Warning {idx} !!:\n{self.coco.cats[idx]['name']} vs {name}")
        # use all classes, cannot use self.metainfo anymore.
        self.cat_ids = self.coco.get_cat_ids()
        self.cat2label = {cat_id: i for i, cat_id in enumerate(self.cat_ids)}
        self.cat_img_map = copy.deepcopy(self.coco.cat_img_map)

        img_ids = self.coco.get_img_ids()
        data_list = []
        total_ann_ids = []
        for img_id in img_ids:
            raw_img_info = self.coco.load_imgs([img_id])[0]
            raw_img_info['img_id'] = img_id

            ann_ids = self.coco.get_ann_ids(img_ids=[img_id])
            raw_ann_info = self.coco.load_anns(ann_ids)
            total_ann_ids.extend(ann_ids)

            parsed_data_info = self.parse_data_info({
                'raw_ann_info':
                    raw_ann_info,
                'raw_img_info':
                    raw_img_info
            })
            data_list.append(parsed_data_info)
        if self.ANN_ID_UNIQUE:
            assert len(set(total_ann_ids)) == len(
                total_ann_ids
            ), f"Annotation ids in '{self.ann_file}' are not unique!"

        del self.coco

        return data_list