import json import os import random from collections import OrderedDict import torch from lib.train.admin import env_settings from lib.train.data import jpeg4py_loader from .base_video_dataset import BaseVideoDataset from .refer import REFER class RefCOCOSeq(BaseVideoDataset): """ The COCO dataset. COCO is an image dataset. Thus, we treat each image as a sequence of length 1. Publication: Microsoft COCO: Common Objects in Context. Tsung-Yi Lin, Michael Maire, Serge J. Belongie, Lubomir D. Bourdev, Ross B. Girshick, James Hays, Pietro Perona, Deva Ramanan, Piotr Dollar and C. Lawrence Zitnick ECCV, 2014 https://arxiv.org/pdf/1405.0312.pdf Download the images along with annotations from http://cocodataset.org/#download. The root folder should be organized as follows. - coco_root - annotations - instances_train2014.json - instances_train2017.json - images - train2014 - train2017 Note: You also have to install the coco pythonAPI from https://github.com/cocodataset/cocoapi. """ def __init__(self, root=None, image_loader=jpeg4py_loader, data_fraction=None, split="train", version="2014", name="refcoco", splitBy="google"): """ args: root - path to the coco dataset. image_loader (default_image_loader) - The function to read the images. If installed, jpeg4py (https://github.com/ajkxyz/jpeg4py) is used by default. Else, opencv's imread is used. data_fraction (None) - Fraction of images to be used. The images are selected randomly. If None, all the images will be used split - 'train' or 'val'. version - version of coco dataset (2014 or 2017) """ root = env_settings().coco_dir if root is None else root root = os.path.join(root,"other") super().__init__('RefCOCOSeq', root, image_loader) self.split = split self.img_pth = os.path.join(root, 'images/mscoco/images/{}{}'.format("train", version)) self.anno_path = os.path.join(root, '{}/instances.json'.format(name)) self.dataset_name = name # Load the COCO set. self.coco_set = REFER(root, dataset=name, splitBy=splitBy) self.cats = self.coco_set.Cats self.class_list = self.get_class_list() self.sequence_list = self._get_sequence_list() if data_fraction is not None: self.sequence_list = random.sample(self.sequence_list, int(len(self.sequence_list) * data_fraction)) self.seq_per_class = self._build_seq_per_class() # get_the_subject_infor self.subject_infor_path = "./" while "resource" not in os.listdir(self.subject_infor_path): self.subject_infor_path = os.path.join(self.subject_infor_path, "../") self.subject_infor_path = os.path.join(self.subject_infor_path, "resource/text_infor/refcoco_4o_v3_extract_mask.json") with open(self.subject_infor_path, 'r') as file: self.subject_infor = json.load(file) def _get_sequence_list(self): ref_list = list(self.coco_set.getRefIds(split=self.split)) seq_list = [a for a in ref_list if self.coco_set.refToAnn[a]['iscrowd'] == 0] return seq_list def is_video_sequence(self): return False def get_num_classes(self): return len(self.class_list) def get_name(self): return self.dataset_name def has_class_info(self): return True def get_class_list(self): class_list = [] for cat_id in self.cats.keys(): class_list.append(self.cats[cat_id]) return class_list def has_segmentation_info(self): return True def get_num_sequences(self): return len(self.sequence_list) def _build_seq_per_class(self): seq_per_class = {} for i, seq in enumerate(self.sequence_list): class_name = self.cats[self.coco_set.refToAnn[seq]['category_id']] if class_name not in seq_per_class: seq_per_class[class_name] = [i] else: seq_per_class[class_name].append(i) return seq_per_class def get_sequences_in_class(self, class_name): return self.seq_per_class[class_name] def get_sequence_info(self, seq_id): anno = self._get_anno(seq_id) bbox = torch.Tensor(anno['bbox']).view(1, 4) # mask = torch.Tensor(self.coco_set.annToMask(anno)).unsqueeze(dim=0) '''2021.1.3 To avoid too small bounding boxes. Here we change the threshold to 50 pixels''' valid = (bbox[:, 2] > 50) & (bbox[:, 3] > 50) visible = valid.clone().byte() nlp = self._read_nlp(seq_id) return {'bbox': bbox, 'valid': valid, 'visible': visible, 'nlp': nlp} def _read_nlp(self, seq_id): ref = self.coco_set.Refs[self.sequence_list[seq_id]] sent = ref['sentences'][-1]['sent'] return sent def _get_anno(self, seq_id): anno = self.coco_set.refToAnn[self.sequence_list[seq_id]] return anno def _get_frames(self, seq_id): path = self.coco_set.loadImgs([self.coco_set.refToAnn[self.sequence_list[seq_id]]['image_id']])[0]['file_name'] img = self.image_loader(os.path.join(self.img_pth, path)) return img def get_meta_info(self, seq_id): try: cat_dict_current = self.cats[self.coco_set.refToAnn[self.sequence_list[seq_id]]['category_id']] object_meta = OrderedDict({'object_class_name': cat_dict_current['name'], 'motion_class': None, 'major_class': cat_dict_current['supercategory'], 'root_class': None, 'motion_adverb': None}) except: object_meta = OrderedDict({'object_class_name': None, 'motion_class': None, 'major_class': None, 'root_class': None, 'motion_adverb': None}) return object_meta def get_class_name(self, seq_id): cat_dict_current = self.cats[self.coco_set.refToAnn[self.sequence_list[seq_id]]['category_id']] return cat_dict_current['name'] def get_subject_mask_infor(self,seq_id): # seq_name = self.sequence_list[seq_id] seq_sub_mask_infor = self.subject_infor[str(seq_id)]["subject_extrack_mask_infor"] return seq_sub_mask_infor def get_frames(self, seq_id=None, frame_ids=None, anno=None): # COCO is an image dataset. Thus we replicate the image denoted by seq_id len(frame_ids) times, and return a # list containing these replicated images. frame = self._get_frames(seq_id) seq_name = self.sequence_list[seq_id] frame_list = [frame.copy() for _ in frame_ids] if anno is None: anno = self.get_sequence_info(seq_id) anno_frames = {} for key, value in anno.items(): if key == 'nlp': anno_frames[key] = [value for _ in frame_ids] else: anno_frames[key] = [value[0, ...] for _ in frame_ids] object_meta = self.get_meta_info(seq_id) subject_mask_infor = self.get_subject_mask_infor(seq_id) if "NoneNone" in subject_mask_infor: subject_mask_infor = [-1] nlp_with_mask = anno_frames["nlp"][0] + "+" + ",".join(map(str, subject_mask_infor)) anno_frames["nlp"] = [nlp_with_mask] # anno_frames["subject_mask_infor"] = subject_mask_infor return frame_list, anno_frames, object_meta def get_path(self, seq_id, frame_ids): img_name = self.coco_set.loadImgs([self.coco_set.refToAnn[self.sequence_list[seq_id]]['image_id']])[0]['file_name'] return [os.path.join(self.img_pth, img_name) for _ in range(len(frame_ids))] def get_ref_id(self, seq_id): return self.sequence_list[seq_id]