def __getitem__(self, idx): img, anno = super(COCODataset, self).__getitem__(idx) # filter crowd annotations # TODO might be better to add an extra field anno = [obj for obj in anno if obj["iscrowd"] == 0] boxes = [obj["bbox"] for obj in anno] boxes = torch.as_tensor(boxes).reshape(-1, 4) # guard against no boxes target = BoxList(boxes, img.size, mode="xywh").convert("xyxy") classes = [obj["category_id"] for obj in anno] classes = [self.json_category_id_to_contiguous_id[c] for c in classes] classes = torch.tensor(classes) target.add_field("labels", classes) masks = [obj["segmentation"] for obj in anno] masks = SegmentationMask(masks, img.size, mode='poly') target.add_field("masks", masks) if anno and "keypoints" in anno[0]: keypoints = [obj["keypoints"] for obj in anno] keypoints = PersonKeypoints(keypoints, img.size) target.add_field("keypoints", keypoints) target = target.clip_to_image(remove_empty=True) if self._transforms is not None: img, target = self._transforms(img, target) return img, target, idx
def __init__(self, method_name='runTest'): super(TestSegmentationMask, self).__init__(method_name) poly = [[ [ 423.0, 306.5, 406.5, 277.0, 400.0, 271.5, 389.5, 277.0, 387.5, 292.0, 384.5, 295.0, 374.5, 220.0, 378.5, 210.0, 391.0, 200.5, 404.0, 199.5, 414.0, 203.5, 425.5, 221.0, 438.5, 297.0, 423.0, 306.5 ], [100, 100, 200, 100, 200, 200, 100, 200], ]] width = 640 height = 480 size = width, height self.P = SegmentationMask(poly, size, 'poly') self.M = SegmentationMask(poly, size, 'poly').convert('mask')
def __call__(self, tensor, target): if self.ratio >= 1.0: return tensor, target self.img_pool.append({'tensor': tensor, 'target': target}) if len(self.img_pool) > self.img_pool_size: self.img_pool.pop(0) c, h, w = tensor.shape if self.size_divisible > 0: h = int(math.ceil(h / self.size_divisible) * self.size_divisible) w = int(math.ceil(w / self.size_divisible) * self.size_divisible) new_h, new_w = int(self.ratio * h), int(self.ratio * w) in_tensor, in_target = scale_jitter(tensor, target, (new_h, new_w)) if len(self.img_pool) < 4: tensor_out, target_out = in_tensor, in_target else: pad_imgs = random.sample(self.img_pool, 3) pad_tensors, pad_targets = [], [] for img in pad_imgs: pad_tensor, pad_target = scale_jitter(img['tensor'], img['target'], (new_h, new_w)) pad_tensors.append(pad_tensor) pad_targets.append(pad_target) crop_boxes = [(0, 0, w - new_w, new_h), (0, 0, new_w, h - new_h), (0, 0, w - new_w, h - new_h)] tensor_out = in_tensor.new(*(c, h, w)).zero_() tensor_out[:c, :new_h, :new_w].copy_(in_tensor) tensor_out[:c, :new_h, new_w:].copy_( pad_tensors[0][:c, :crop_boxes[0][3], :crop_boxes[0][2]]) tensor_out[:c, new_h:, :new_w].copy_( pad_tensors[1][:c, :crop_boxes[1][3], :crop_boxes[1][2]]) tensor_out[:c, new_h:, new_w:].copy_( pad_tensors[2][:c, :crop_boxes[2][3], :crop_boxes[2][2]]) crop_targets = [] for i, pad_target in enumerate(pad_targets): crop_target = pad_target.crop(crop_boxes[i]) ious = crop_target.area() / pad_target.area() crop_target = crop_target[ious >= self.iou_threshold] crop_targets.append(crop_target) offsets_box = [ torch.Tensor([0.0, 0.0, 0.0, 0.0]), torch.Tensor([new_w, 0.0, new_w, 0.0]), torch.Tensor([0.0, new_h, 0.0, new_h]), torch.Tensor([new_w, new_h, new_w, new_h]) ] target_out = torch.cat([ target.bbox + offsets_box[i] for i, target in enumerate([in_target] + crop_targets) ], dim=0) target_out = BoxList(target_out, (w, h), mode='xyxy') target_out.add_field( 'labels', torch.cat([ target.extra_fields['labels'] for target in ([in_target] + crop_targets) ], dim=-1)) polys_list = [[ poly.polygons[0] for poly in target.extra_fields['masks'].instances.polygons ] for target in ([in_target] + crop_targets)] offsets_mask = [[0.0, 0.0], [new_w, 0.0], [0.0, new_h], [new_w, new_h]] syn_mask = [] for i, polys in enumerate(polys_list): syn_mask += [[ list( np.array(poly) + np.array(offsets_mask[i] * int(len(poly) / 2))) ] for poly in polys] syn_mask = SegmentationMask(syn_mask, (w, h), mode='poly') target_out.add_field('masks', syn_mask) return tensor_out, target_out
def __getitem__(self, idx): img, anno = super(COCODataset, self).__getitem__(idx) # filter crowd annotations # TODO might be better to add an extra field # anno = [obj for obj in anno if obj["iscrowd"] == 0] dict_key = self.ids[idx] diags = self.diag_dict[()][dict_key][0] invalids = self.diag_dict[()][dict_key][1] centers = self.diag_dict[()][dict_key][2] if not img.getbbox(): img_data = img.convert("RGB") print(img_data) raise Exception('Image Error for key : {}'.format(dict_key)) boxes = [ obj["bbox"] for indx, obj in enumerate(anno) if indx not in invalids ] boxes = torch.as_tensor(boxes).reshape(-1, 4) # guard against no boxes target = BoxList(boxes, img.size, mode="xywh").convert("xyxy") classes = [ obj["category_id"] for indx, obj in enumerate(anno) if indx not in invalids ] classes = [self.json_category_id_to_contiguous_id[c] for c in classes] classes = torch.tensor(classes) target.add_field("labels", classes) masks = [obj["segmentation"] for obj in anno] masks = SegmentationMask(masks, img.size, mode='poly') target.add_field("masks", masks) diags = torch.as_tensor(diags, dtype=torch.float).reshape(-1, 16) target.add_field("diagonals", diags) centers = torch.as_tensor(centers, dtype=torch.float).reshape(-1, 2) target.add_field("centers", centers) # img_id = self.id_to_img_map[idx] # img_data = self.coco.imgs[img_id] assert len(diags) == len( boxes ), 'Mismatch in boxes and diags for id: {} with diag size {} and bbox size {} \n and boxes are: {}'.format( dict_key, diags.size(), boxes.size(), boxes) # coco_visualization(target, img_data, pre_train=True) if anno and "keypoints" in anno[0]: keypoints = [obj["keypoints"] for obj in anno] keypoints = PersonKeypoints(keypoints, img.size) target.add_field("keypoints", keypoints) #target = target.clip_to_image(remove_empty=True) # if self._transforms is not None: # img, target = self._transforms(img, target) return img, target, idx
classes = [obj["category_id"] for obj in anno] print('classes', classes) # from pycocotools.coco import COCO # coco = COCO(annFile) # json_category_id_to_contiguous_id = { # v: i + 1 for i, v in enumerate(coco.getCatIds())} # classes = [json_category_id_to_contiguous_id[c] for c in classes] # print('classes', classes) classes = 7 classes = torch.tensor(classes) target.add_field("labels", classes) #self.extra_fields["labels"] = 7 masks = [obj["segmentation"] for obj in anno] print('masks', masks) masks = SegmentationMask(masks, img.size, mode='poly') print('masks', masks) target.add_field("masks", masks) #self.extra_fields["masks"] = SegmentationMask(num_instances=1, image_width=640, image_height=427, mode=poly) print(target) print(target.bbox) #target = target.clip_to_image(remove_empty=True) transforms = build_transforms(cfg, True) img, target = transforms(img, target) print(img.shape) print(target) print(target.bbox) #model = build_detection_model(cfg) #loss_dict = model(img, target) #print(loss_dict) labels_per_im = target.get_field("labels")
def _geometric_aug_func(x, target, angle=0, translate=(0, 0), scale=1, shear=(0, 0), hflip=False, boxes_sample_prob=[], scale_ratio=1.0): boxes_and_labels = [(target.bbox[i], target.extra_fields['labels'][i], target.extra_fields['masks'].instances.polygons[i]) for i in range(len(target.bbox)) if random.random() < boxes_sample_prob[i]] boxes = [b_and_l[0] for b_and_l in boxes_and_labels] labels = [b_and_l[1] for b_and_l in boxes_and_labels] polys = [b_and_l[2].polygons[0] for b_and_l in boxes_and_labels] if random.random() < 0.5: angle *= -1 translate = (-translate[0], -translate[1]) shear = (-shear[0], -shear[1]) translate = (0, 0) height, width = x.shape[1], x.shape[2] x_crops = [] boxes_crops = [] boxes_new = [] polys_new = [] labels_new = [] for i, box in enumerate(boxes): box_crop = scale_area(box, height, width, scale_ratio) y1, x1, y2, x2 = box_crop.long() x_crop = x[:, x1:x2, y1:y2] boxes_crops.append(box_crop) if x1 >= x2 or y1 >= y2: x_crops.append(x_crop) continue if hflip: x_crop = x_crop.flip(-1) elif translate[0] + translate[1] != 0: offset_y = (y2 + translate[0]).clamp(0, width).long().tolist() - y2 offset_x = (x2 + translate[1]).clamp(0, height).long().tolist() - x2 if offset_x != 0 or offset_y != 0: offset = [offset_y, offset_x] boxes_new.append(box + torch.Tensor(offset * 2)) polys_new.append( np.array(polys[i]) + np.array(offset * int(len(polys[i]) / 2))) labels_new.append(labels[i]) else: x_crop = transforms.functional.to_pil_image(x_crop.cpu()) x_crop = transforms.functional.affine( x_crop, angle, translate, scale, shear, resample=2, fillcolor=tuple([int(i) for i in pixel_mean])) x_crop = transforms.functional.to_tensor(x_crop).to(x.device) x_crops.append(x_crop) y = _transform(x, x_crops, boxes_crops, translate) if translate[0] + translate[1] != 0 and len(boxes_new) > 0: target.bbox = torch.cat((target.bbox, torch.stack(boxes_new))) target.extra_fields['labels'] = torch.cat( (target.extra_fields['labels'], torch.Tensor(labels_new).long())) polys = [ poly.polygons[0] for poly in target.extra_fields['masks'].instances.polygons ] + polys_new target.extra_fields['masks'] = SegmentationMask( [[list(np.array(poly))] for poly in polys], target.extra_fields['masks'].size, mode='poly') return y, target