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