MLVLab/Human_Object_Interaction
1
1# Copyright (c) Kakaobrain, Inc. and its affiliates. All Rights Reserved2"""3V-COCO dataset which returns image_id for evaluation.4"""5from pathlib import Path6 7from PIL import Image8import os9import numpy as np10import json11import torch12import torch.utils.data13import torchvision14 15from torch.utils.data import Dataset16from pycocotools.coco import COCO17from pycocotools import mask as coco_mask18 19from hotr.data.datasets import builtin_meta20import hotr.data.transforms.transforms as T21 22class VCocoDetection(Dataset):23 def __init__(self,24 img_folder,25 ann_file,26 all_file,27 filter_empty_gt=True,28 transforms=None):29 self.img_folder = img_folder30 self.file_meta = dict()31 self._transforms = transforms32 33 self.ann_file = ann_file34 self.all_file = all_file35 self.filter_empty_gt = filter_empty_gt36 37 # COCO initialize38 self.coco = COCO(self.all_file)39 self.COCO_CLASSES = builtin_meta._get_coco_instances_meta()['coco_classes']40 self.file_meta['coco_classes'] = self.COCO_CLASSES41 42 # Load V-COCO Dataset43 self.vcoco_all = self.load_vcoco(self.ann_file)44 45 # Save COCO annotation data46 self.image_ids = sorted(list(set(self.vcoco_all[0]['image_id'].reshape(-1))))47 48 # Filter Data49 if filter_empty_gt:50 self.filter_image_id()51 self.img_infos = self.load_annotations()52 53 # Refine Data54 self.save_action_name()55 self.mapping_inst_action_to_action()56 self.load_subobj_classes()57 self.CLASSES = self.act_list58 59 ############################################################################60 # Load V-COCO Dataset61 ############################################################################62 def load_vcoco(self, dir_name=None):63 with open(dir_name, 'rt') as f:64 vsrl_data = json.load(f)65 66 for i in range(len(vsrl_data)):67 vsrl_data[i]['role_object_id'] = np.array(vsrl_data[i]['role_object_id']).reshape((len(vsrl_data[i]['role_name']),-1)).T68 for j in ['ann_id', 'label', 'image_id']:69 vsrl_data[i][j] = np.array(vsrl_data[i][j]).reshape((-1,1))70 71 return vsrl_data72 73 ############################################################################74 # Refine Data75 ############################################################################76 def save_action_name(self):77 self.inst_act_list = list()78 self.act_list = list()79 80 # add instance action human classes81 self.num_subject_act = 082 for vcoco in self.vcoco_all:83 self.inst_act_list.append('human_' + vcoco['action_name'])84 self.num_subject_act += 185 86 # add instance action object classes87 for vcoco in self.vcoco_all:88 if len(vcoco['role_name']) == 3:89 self.inst_act_list.append('object_' + vcoco['action_name']+'_'+vcoco['role_name'][1])90 self.inst_act_list.append('object_' + vcoco['action_name']+'_'+vcoco['role_name'][2])91 elif len(vcoco['role_name']) < 2:92 continue93 else:94 self.inst_act_list.append('object_' + vcoco['action_name']+'_'+vcoco['role_name'][-1]) # when only two roles95 96 # add action classes97 for vcoco in self.vcoco_all:98 if len(vcoco['role_name']) == 3:99 self.act_list.append(vcoco['action_name']+'_'+vcoco['role_name'][1])100 self.act_list.append(vcoco['action_name']+'_'+vcoco['role_name'][2])101 else:102 self.act_list.append(vcoco['action_name']+'_'+vcoco['role_name'][-1])103 104 # add to meta105 self.file_meta['action_classes'] = self.act_list106 107 def mapping_inst_action_to_action(self):108 sub_idx = 0109 obj_idx = self.num_subject_act110 111 self.sub_label_to_action = list()112 self.obj_label_to_action = list()113 114 for vcoco in self.vcoco_all:115 role_name = vcoco['role_name']116 117 self.sub_label_to_action.append(sub_idx)118 if len(role_name) == 3 :119 self.sub_label_to_action.append(sub_idx)120 self.obj_label_to_action.append(obj_idx)121 self.obj_label_to_action.append(obj_idx+1)122 obj_idx += 2123 elif len(role_name) == 2:124 self.obj_label_to_action.append(obj_idx)125 obj_idx += 1126 else:127 self.obj_label_to_action.append(0)128 129 sub_idx += 1130 131 def load_subobj_classes(self):132 self.vcoco_labels = dict()133 for img in self.image_ids:134 self.vcoco_labels[img] = dict()135 self.vcoco_labels[img]['boxes'] = np.empty((0, 4), dtype=np.float32)136 self.vcoco_labels[img]['categories'] = np.empty((0), dtype=np.int32)137 138 ann_ids = self.coco.getAnnIds(imgIds=img, iscrowd=None)139 objs = self.coco.loadAnns(ann_ids)140 141 valid_ann_ids = []142 143 for i, obj in enumerate(objs):144 if 'ignore' in obj and obj['ignore'] == 1: continue145 146 x1 = obj['bbox'][0]147 y1 = obj['bbox'][1]148 x2 = x1 + np.maximum(0., obj['bbox'][2] - 1.)149 y2 = y1 + np.maximum(0., obj['bbox'][3] - 1.)150 151 if obj['area'] > 0 and x2 > x1 and y2 > y1:152 bbox = np.array([x1, y1, x2, y2]).reshape(1, -1)153 cls = obj['category_id']154 self.vcoco_labels[img]['boxes'] = np.concatenate([self.vcoco_labels[img]['boxes'], bbox], axis=0)155 self.vcoco_labels[img]['categories'] = np.concatenate([self.vcoco_labels[img]['categories'], [cls]], axis=0)156 157 valid_ann_ids.append(ann_ids[i])158 159 num_valid_objs = len(valid_ann_ids)160 161 self.vcoco_labels[img]['agent_actions'] = -np.ones((num_valid_objs, self.num_action()), dtype=np.int32)162 self.vcoco_labels[img]['obj_actions'] = np.zeros((num_valid_objs, self.num_action()), dtype=np.int32)163 self.vcoco_labels[img]['role_id'] = -np.ones((num_valid_objs, self.num_action()), dtype=np.int32)164 165 for ix, ann_id in enumerate(valid_ann_ids):166 in_vcoco = np.where(self.vcoco_all[0]['ann_id'] == ann_id)[0]167 if in_vcoco.size > 0:168 self.vcoco_labels[img]['agent_actions'][ix, :] = 0169 170 agent_act_id = 0171 obj_act_id = -1172 for i, x in enumerate(self.vcoco_all):173 has_label = np.where(np.logical_and(x['ann_id'] == ann_id, x['label'] == 1))[0]174 if has_label.size > 0:175 assert has_label.size == 1176 rids = x['role_object_id'][has_label]177 178 if rids.shape[1] == 3:179 self.vcoco_labels[img]['agent_actions'][ix, agent_act_id] = 1180 self.vcoco_labels[img]['agent_actions'][ix, agent_act_id+1] = 1181 agent_act_id += 2182 else:183 self.vcoco_labels[img]['agent_actions'][ix, agent_act_id] = 1184 agent_act_id += 1185 if rids.shape[1] == 1 : obj_act_id += 1186 187 for j in range(1, rids.shape[1]):188 obj_act_id += 1189 if rids[0, j] == 0: continue # no role190 aid = np.where(valid_ann_ids == rids[0, j])[0]191 192 self.vcoco_labels[img]['role_id'][ix, obj_act_id] = aid193 self.vcoco_labels[img]['obj_actions'][aid, obj_act_id] = 1194 195 else:196 rids = x['role_object_id'][0]197 if rids.shape[0] == 3:198 agent_act_id += 2199 obj_act_id += 2200 else:201 agent_act_id += 1202 obj_act_id += 1203 204 ############################################################################205 # Annotation Loader206 ############################################################################207 # >>> 1. instance208 def load_instance_annotations(self, image_index):209 num_ann = self.vcoco_labels[image_index]['boxes'].shape[0]210 inst_action = np.zeros((num_ann, self.num_inst_action()), np.int)211 inst_bbox = np.zeros((num_ann, 4), dtype=np.float32)212 inst_category = np.zeros((num_ann, ), dtype=np.int)213 214 for idx in range(num_ann):215 inst_bbox[idx] = self.vcoco_labels[image_index]['boxes'][idx]216 inst_category[idx]= self.vcoco_labels[image_index]['categories'][idx] #+ 1 # category 1 ~ 81217 218 if inst_category[idx] == 1:219 act = self.vcoco_labels[image_index]['agent_actions'][idx]220 inst_action[idx, :self.num_subject_act] = act[np.unique(self.sub_label_to_action, return_index=True)[1]]221 222 # when person is the obj223 act = self.vcoco_labels[image_index]['obj_actions'][idx] # when person is the obj224 if act.any():225 inst_action[idx, self.num_subject_act:] = act[np.nonzero(self.obj_label_to_action)[0]]226 if inst_action[idx, :self.num_subject_act].sum(axis=-1) < 0:227 inst_action[idx, :self.num_subject_act] = 0228 else:229 act = self.vcoco_labels[image_index]['obj_actions'][idx]230 inst_action[idx, self.num_subject_act:] = act[np.nonzero(self.obj_label_to_action)[0]]231 232 # >>> For Objects that are in COCO but not in V-COCO,233 # >>> Human -> [-1 * 26, 0 * 25]234 # >>> Object -> [0 * 51]235 # >>> Don't return anything for actions with max 0 or max -1236 max_val = inst_action.max(axis=1)237 if (max_val > 0).sum() == 0:238 print(f"No Annotations for {image_index}")239 print(inst_action)240 print(self.vcoco_labels[image_index]['agent_actions'][idx])241 print(self.vcoco_labels[image_index]['obj_actions'][idx])242 243 return inst_bbox[max_val > 0], inst_category[max_val > 0], inst_action[max_val > 0]244 245 # >>> 2. pair246 def load_pair_annotations(self, image_index):247 num_ann = self.vcoco_labels[image_index]['boxes'].shape[0]248 pair_action = np.zeros((0, self.num_action()), np.int)249 pair_bbox = np.zeros((0, 8), dtype=np.float32)250 pair_target = np.zeros((0, ), dtype=np.int)251 252 for idx in range(num_ann):253 h_box = self.vcoco_labels[image_index]['boxes'][idx]254 h_cat = self.vcoco_labels[image_index]['categories'][idx]255 if h_cat != 1 : continue # human_id = 1256 257 h_act = self.vcoco_labels[image_index]['agent_actions'][idx]258 if np.any((h_act==-1)) : continue259 260 o_act = dict()261 for aid in range(self.num_action()):262 if h_act[aid] == 0 : continue263 o_id = self.vcoco_labels[image_index]['role_id'][idx, aid]264 if o_id not in o_act : o_act[o_id] = list()265 o_act[o_id].append(aid)266 267 for o_id in o_act.keys():268 if o_id == -1:269 o_box = -np.ones((4, ))270 o_cat = -1 # target is background271 else:272 o_box = self.vcoco_labels[image_index]['boxes'][o_id]273 o_cat = self.vcoco_labels[image_index]['categories'][o_id] # category 0 ~ 80274 275 box = np.concatenate([h_box, o_box]).astype(np.float32)276 act = np.zeros((1, self.num_action()), np.int)277 tar = np.zeros((1, ), np.int)278 tar[0] = o_cat #+ 1 # category 1 ~ 81279 for o_aid in o_act[o_id] : act[0, o_aid] = 1280 281 pair_action = np.concatenate([pair_action, act], axis=0)282 pair_bbox = np.concatenate([pair_bbox, np.expand_dims(box, axis=0)], axis=0)283 pair_target = np.concatenate([pair_target, tar], axis=0)284 285 return pair_bbox, pair_action, pair_target286 287 # >>> 3. image infos288 def load_annotations(self):289 img_infos = []290 for i in self.image_ids:291 info = self.coco.loadImgs([i])[0]292 img_infos.append(info)293 return img_infos294 295 ############################################################################296 # Check Method297 ############################################################################298 def sum_action_ann_for_id(self, find_idx):299 sum = 0300 for action_ann in self.vcoco_all:301 img_ids = action_ann['image_id']302 img_labels = action_ann['label']303 304 final_inds = img_ids[img_labels == 1]305 306 if (find_idx in final_inds):307 sum += 1308 # sum of class-wise existence309 return (sum > 0)310 311 def filter_image_id(self):312 empty_gt_list = []313 for img_id in self.image_ids:314 if not self.sum_action_ann_for_id(img_id):315 empty_gt_list.append(img_id)316 317 for remove_id in empty_gt_list:318 rm_idx = self.image_ids.index(remove_id)319 self.image_ids.remove(remove_id)320 321 ############################################################################322 # Preprocessing323 ############################################################################324 def prepare_img(self, idx):325 img_info = self.img_infos[idx]326 image = Image.open(os.path.join(self.img_folder, img_info['file_name'])).convert('RGB')327 target = self.get_ann_info(idx)328 329 w, h = image.size330 target["orig_size"] = torch.as_tensor([int(h), int(w)])331 target["size"] = torch.as_tensor([int(h), int(w)])332 333 if self._transforms is not None:334 img, target = self._transforms(image, target) # "size" gets converted here335 336 return img, target337 338 ############################################################################339 # Get Method340 ############################################################################341 def __getitem__(self, idx):342 img, target = self.prepare_img(idx)343 return img, target344 345 def __len__(self):346 return len(self.image_ids)347 348 def get_human_label_idx(self):349 return self.sub_label_to_action350 351 def get_object_label_idx(self):352 return self.obj_label_to_action353 354 def get_image_ids(self):355 return self.image_ids356 357 def get_categories(self):358 return self.COCO_CLASSES359 360 def get_inst_action(self):361 return self.inst_act_list362 363 def get_actions(self):364 return self.act_list365 366 def get_human_action(self):367 return self.inst_act_list[:self.num_subject_act]368 369 def get_object_action(self):370 return self.inst_act_list[self.num_subject_act:]371 372 def get_ann_info(self, idx):373 img_idx = int(self.image_ids[idx])374 375 # load each annotation376 inst_bbox, inst_label, inst_actions = self.load_instance_annotations(img_idx)377 pair_bbox, pair_actions, pair_targets = self.load_pair_annotations(img_idx)378 379 sample = {380 'image_id' : torch.tensor([img_idx]),381 'boxes': torch.as_tensor(inst_bbox, dtype=torch.float32),382 'labels': torch.tensor(inst_label, dtype=torch.int64),383 'inst_actions': torch.tensor(inst_actions, dtype=torch.int64),384 'pair_boxes': torch.as_tensor(pair_bbox, dtype=torch.float32),385 'pair_actions': torch.tensor(pair_actions, dtype=torch.int64),386 'pair_targets': torch.tensor(pair_targets, dtype=torch.int64),387 }388 389 return sample390 391 ############################################################################392 # Number Method393 ############################################################################394 def num_category(self):395 return len(self.COCO_CLASSES)396 397 def num_action(self):398 return len(self.act_list)399 400 def num_inst_action(self):401 return len(self.inst_act_list)402 403 def num_human_act(self):404 return len(self.inst_act_list[:self.num_subject_act])405 406 def num_object_act(self):407 return len(self.inst_act_list[self.num_subject_act:])408 409def make_hoi_transforms(image_set):410 normalize = T.Compose([411 T.ToTensor(),412 T.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])413 ])414 415 scales = [480, 512, 544, 576, 608, 640, 672, 704, 736, 768, 800]416 417 if image_set == 'train':418 return T.Compose([419 T.RandomHorizontalFlip(),420 T.ColorJitter(.4, .4, .4),421 T.RandomSelect(422 T.RandomResize(scales, max_size=1333),423 T.Compose([424 T.RandomResize([400, 500, 600]),425 T.RandomSizeCrop(384, 600),426 T.RandomResize(scales, max_size=1333),427 ])428 ),429 normalize,430 ])431 432 if image_set == 'val':433 return T.Compose([434 T.RandomResize([800], max_size=1333),435 normalize,436 ])437 438 if image_set == 'test':439 return T.Compose([440 T.RandomResize([800], max_size=1333),441 normalize,442 ])443 444 raise ValueError(f'unknown {image_set}')445 446def build(image_set, args):447 root = Path(args.data_path)448 assert root.exists(), f'provided V-COCO path {root} does not exist'449 PATHS = {450 "train": (root / "coco/images/train2014/", root / "data/vcoco" / 'vcoco_trainval.json'),451 "val": (root / "coco/images/val2014", root / "data/vcoco" / 'vcoco_test.json'),452 "test": (root / "coco/images/val2014", root / "data/vcoco" / 'vcoco_test.json'),453 }454 455 img_folder, ann_file = PATHS[image_set]456 all_file = root / "data/instances_vcoco_all_2014.json"457 dataset = VCocoDetection(458 img_folder = img_folder,459 ann_file = ann_file,460 all_file = all_file,461 filter_empty_gt=True,462 transforms = make_hoi_transforms(image_set)463 )464 dataset.file_meta['dataset_file'] = args.dataset_file465 dataset.file_meta['image_set'] = image_set466 467 return dataset468