def __getitem__(self, idx):
"""Generate an image from the specs of the given image ID.
Typically this function loads the image from a file, but
in this case it generates the image on the fly from the
specs in image_info.
"""
image = Image.fromarray(self.load_image(idx))
segmentation_mask, masks, labels, boxes = self.load_mask(idx)
# create a BoxList from the boxes
boxlist = BoxList(boxes, image.size, mode="xyxy")
# add the labels to the boxlist
boxlist.add_field("labels", torch.tensor(labels))
# Add masks to the boxlist
masks = np.transpose(masks, (2, 0, 1))
masks = SegmentationMask(torch.tensor(masks), image.size, "mask")
boxlist.add_field("masks", masks)
# Add semantic segmentation masks to the boxlist for panoptic segmentation
segmentation_mask = np.transpose(segmentation_mask, (2, 0, 1))
seg_masks = SegmentationMask(torch.tensor(segmentation_mask),
image.size, "mask")
boxlist.add_field("seg_masks", seg_masks)
# Important line! dont forget to add this
if self.transforms:
image, boxlist = self.transforms(image, boxlist)
# return the image, the boxlist and the idx in your dataset
return image, boxlist, idx
def viz_coco_example(img, imginfo, annotations, classnames):
height, width = imginfo['height'], imginfo['width']
assert (img.shape[0] == height)
assert (img.shape[1] == width)
masks = SegmentationMask([ann['segmentation'] for ann in annotations],
size=(width, height),
mode='poly')
bboxes = []
classids = []
for ann in annotations:
x, y, width, height = ann['bbox']
x1, y1, x2, y2 = x, y, x + width, y + height
bboxes.append([y1, x1, y2, x2])
classids.append(ann['category_id'] -
1) # category_ids are 1-indexed in COCO datasets
classids = np.array(list(map(int, classids)))
bboxes = np.stack(bboxes, axis=0)
masks = masks.get_mask_tensor().numpy()
if masks.ndim == 2:
masks = masks[np.newaxis, :, :]
masks = np.transpose(masks, (1, 2, 0))
display_instances(img,
boxes=bboxes,
masks=masks,
class_ids=classids,
class_names=classnames)
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 get_groundtruth(self, idx: int):
bboxes = []
masks = []
svg_name = self.data_dir / "svg" / "{:s}.svg".format(self.ids[idx])
with svg_name.open('r') as f_svg:
svg = f_svg.read()
svg_xml = et.fromstring(svg)
for i in range(len(svg_xml[0])):
svg_xml = et.fromstring(svg)
svg_xml[0][0] = svg_xml[0][i]
del svg_xml[0][1:]
svg_one = et.tostring(svg_xml, method='xml')
y_png = cairosvg.svg2png(bytestring=svg_one)
y_img = Image.open(io.BytesIO(y_png))
mask = (np.array(y_img)[:, :, 3] > 0)
assert mask.shape == (self.size, self.size), (self.ids[idx])
bboxes.append(self.get_bbox(mask, "xyxy"))
masks.append(mask.astype(np.uint8))
image_size_dict = self.get_img_info(idx)
image_size = (image_size_dict["height"], image_size_dict["width"])
boxlist = BoxList(bboxes, image_size, mode="xyxy")
boxlist.add_field("labels", torch.tensor([1] * len(bboxes)))
boxlist.add_field("mask", SegmentationMask(masks, image_size, "mask"))
return self.svg_to_png(svg), boxlist
def __getitem__(self, idx):
img_name = os.path.join(self.data_dir_img, self.img_name[idx])
image = open_img(img_name).convert('RGB')
#print(self.size, image.size)
i, j, h, w = self.get_params(image, self.size)
#print("image size : {} ; {}, {}, {}, {}".format(image.size, i, j, h, w))
image = F.crop(image, i, j, h, w)
boxes = torch.as_tensor(self.gt[idx]).reshape(
-1, 4) # guard against no boxes
#print(boxes, self.gt[idx])
#print("Boxes :", boxes.shape)
boxes = boxes - torch.as_tensor([i, j, 0, 0])
#print(boxes)
target = BoxList(boxes, image.size, mode="xywh").convert("xyxy")
classes = torch.as_tensor(self.cls_idx[idx]).reshape(-1)
#print("Classes :", classes.shape)
target.add_field("labels", classes)
masks = torch.as_tensor(1 * self.gen_mask_from_img(image))
masks = SegmentationMask(masks, image.size, mode='mask')
target.add_field("masks", masks)
target = target.clip_to_image(remove_empty=True)
if self._transform:
image, target = self._transform(image, target)
return image, target, idx
def __getitem__(self, item):
im_name = os.path.basename(self.image_lists[item])
#print(self.image_lists[item])
img = Image.open(self.image_lists[item]).convert("RGB")
width, height = img.size
if self.gts_dir is not None:
gt_path = os.path.join(self.gts_dir, im_name + '.txt')
#ipdb.set_trace()
words, boxes, charsbbs, segmentations = self.load_gt_from_txt(
gt_path, height, width)
if words[0] == '':
use_char_ann = False
else:
use_char_ann = True
if not self.use_charann:
use_char_ann = False
target = BoxList(boxes[:, :4],
img.size,
mode="xyxy",
use_char_ann=use_char_ann)
classes = torch.ones(len(boxes))
target.add_field("labels", classes)
masks = SegmentationMask(segmentations, img.size)
target.add_field("masks", masks)
char_masks = SegmentationCharMask(charsbbs,
words=words,
use_char_ann=use_char_ann,
size=img.size)
target.add_field("char_masks", char_masks)
else:
target = None
if self.transforms is not None:
img, target = self.transforms(img, target)
if self.vis:
new_im = img.numpy().copy().transpose(
[1, 2, 0]) + [102.9801, 115.9465, 122.7717]
new_im = Image.fromarray(new_im.astype(np.uint8)).convert('RGB')
mask = target.extra_fields['masks'].polygons[0].convert('mask')
mask = Image.fromarray(
(mask.numpy() * 255).astype(np.uint8)).convert('RGB')
if self.use_charann:
m, _ = target.extra_fields['char_masks'].chars_boxes[
0].convert('char_mask')
color = self.creat_color_map(37, 255)
color_map = color[m.numpy().astype(np.uint8)]
char = Image.fromarray(color_map.astype(
np.uint8)).convert('RGB')
char = Image.blend(char, new_im, 0.5)
else:
char = new_im
new = Image.blend(char, mask, 0.5)
img_draw = ImageDraw.Draw(new)
#ipdb.set_trace()
for box in target.bbox.numpy():
box = list(box)
box = box[:2] + [box[2], box[1]] + box[2:] + [box[0], box[3]
] + box[:2]
img_draw.line(box, fill=(255, 0, 0), width=2)
new.save('./vis/char_' + im_name)
return img, target, self.image_lists[item]
def __getitem__(self, idx):
img, anno = super(COCODataset, self).__getitem__(idx)
anno = []
with open(os.path.join(self.root_json, "json", str(self.ids[idx])),
'r') as f:
for line in f:
anno.append(ujson.loads(line))
# 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!!!
# classes = [obj["category_id"] for obj in anno]
classes = [1 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)
target.add_field("masks", masks)
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 __getitem__(self, index):
masks = []
while len(masks) == 0:
item = self.ids[index]
img_id = item['img_id']
winSize = int(self.maxWS * windowSizeRatio[self.split])
if self.split == 'train':
img_path = self.patch_dir[index]
img = Image.open(img_path)
else:
img = self.image_pool[img_id].read_region(
(item['location'][0], item['location'][1]), 0,
(winSize, winSize)).convert("RGB")
masks, bboxes, labels = self.filter_gt(item)
if len(masks) == 0 or (1 not in set(labels)
or 2 not in set(labels)):
print('No mask for index:', index)
index = np.random.choice(len(self.ids))
bboxes = torch.as_tensor(bboxes).reshape(-1,
4) # guard against no boxes
target = BoxList(bboxes, img.size, mode='xyxy')
labels = torch.tensor(labels)
target.add_field("labels", labels)
masks = SegmentationMask(masks, img.size)
target.add_field("masks", masks)
target = target.clip_to_image(remove_empty=True)
img, target = self.transforms(img, target)
return img, target, index
def project_boxes_on_boxes(matched_bboxes, proposals, discretization_size):
masks = []
M = discretization_size
original_size = proposals.size
proposals = proposals.bbox.to(torch.device("cpu"))
matched_bboxes = matched_bboxes.to(torch.device("cpu"))
# Generate segmentation masks based on matched_bboxes
polygons = []
for matched_bbox in matched_bboxes:
x1, y1, x2, y2 = matched_bbox[0], matched_bbox[1], matched_bbox[
2], matched_bbox[3]
p = [[x1, y1, x1, y2, x2, y2, x2, y1]]
polygons.append(p)
segmentation_masks = SegmentationMask(polygons, original_size)
for segmentation_mask, proposal in zip(segmentation_masks, proposals):
# crop the masks, resize them to the desired resolution and
# then convert them to the tensor representation,
# instead of the list representation that was used
cropped_mask = segmentation_mask.crop(proposal)
size = cropped_mask.size
# LVIS dataset could have tiny boxes
if size[0] > 1e-2 and size[1] > 1e-2:
scaled_mask = cropped_mask.resize((M, M))
mask = scaled_mask.convert(mode="mask")
mask_tensor = mask.get_mask_tensor()
else:
logger = logging.getLogger("maskrcnn_benchmark.trainer")
logger.warning("%s is too small", cropped_mask)
mask_tensor = torch.zeros((28, 28), dtype=torch.uint8)
masks.append(mask_tensor)
return torch.stack(masks, dim=0).to(dtype=torch.float32)
def __getitem__(self, idx):
import pdb;pdb.set_trace()
print("you have reached micr datatset get method")
img, anno = super(MICRDataset, 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)
if anno and "segmentation" in anno[0]:
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 __call__(self, image, target=None):
if target is None:
return image
if not self.do:
return image, target
keypoints = target.get_field("keypoints")
polygons = list(
map(lambda x: x.polygons[0].numpy(), keypoints.instances.polygons))
# polygons = np.stack( polygons, axis=0 ).reshape( (-1, 4, 2) )
# rrects = list( map( self._to_rrect, polygons ) )
rrects_np = np.array(rrects, dtype=np.float32).reshape((-1, 8))
xmins = np.min(rrects_np[:, ::2], axis=1)
ymins = np.min(rrects_np[:, 1::2], axis=1)
xmaxs = np.max(rrects_np[:, ::2], axis=1)
ymaxs = np.max(rrects_np[:, 1::2], axis=1)
xyxy = np.vstack([xmins, ymins, xmaxs, ymaxs]).transpose()
boxes = torch.from_numpy(xyxy).reshape(-1, 4) # guard against no boxes
new_target = BoxList(boxes, image.size, mode="xyxy")
new_target._copy_extra_fields(target)
new_keypoints = SegmentationMask(rrects_np.reshape(
(-1, 1, 8)).tolist(),
image.size,
mode='poly')
new_target.add_field("keypoints", new_keypoints)
return image, new_target
def __getitem__(self, idx):
name = self.id_to_img_map[idx]
mask_name = self.img_name2mask_name(name)
img_path = os.path.join(self.root, name)
mask_path = os.path.join(self.mask_root, mask_name)
img = Image.open(img_path)
mask = np.array(Image.open(mask_path))
mask = mask / 255
mask = mask.astype(np.uint8)
# TODO might be better to add an extra field
annos = self.annos[name]
boxes = [annos[:4]]
boxes = torch.as_tensor(boxes).reshape(-1, 4) # guard against no boxes
target = BoxList(boxes, img.size, mode="xyxy")
classes = [annos[4]]
classes = torch.tensor(classes)
target.add_field("labels", classes)
masks = [mask]
masks = SegmentationMask(masks, img.size, mode='mask')
target.add_field("masks", masks)
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 __getitem__(self, idx):
img_idx,anno_idx=self.ids[idx][1],self.ids[idx][0]
img = Image.open(os.path.join(self.root, img_idx)).convert('RGB')
with open(os.path.join(self.root,anno_idx)) as fr:
anno=json.load(fr)
image_height=anno['image_height']
image_width=anno['image_width']
boxes=[]
classes=[]
masks=[]
for each in anno['bboxes']:
x_min=each['x_min']*image_width
x_max=each['x_max']*image_width
y_min=each['y_min']*image_height
y_max=each['y_max']*image_height
mask=each['mask']
if len(mask)==0:
continue
classes.append(each['label'])
boxes.append([x_min,y_min,x_max-x_min,y_max-y_min])
mask=[ [i,v] for i,v in zip(mask['all_points_x'],mask['all_points_y'])]
mask=list(itertools.chain.from_iterable(mask))
masks.append([mask])
boxes = torch.as_tensor(boxes).reshape(-1, 4) # guard against no boxes
target = BoxList(boxes, img.size, mode="xywh").convert("xyxy")
classes=[self.classmap[c] for c in classes]
classes = torch.tensor(classes)
target.add_field("labels", classes)
#masks
masks = SegmentationMask(masks, img.size, mode='poly')
target.add_field("masks", masks)
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 __getitem__(self, idx):
img_name = self.id_to_img_map[idx]
img = utils.pil_load_img(os.path.join(self.root, img_name))
anno = utils.read_anno(self.annotations, img_name)
# filter illegal
anno = [obj for obj in anno if not obj['illegibility']]
# bounding boxes
boxes = [utils.generate_rbox(obj["points"], np.array(img).shape[:2]) for obj in anno]
boxes = torch.as_tensor(boxes).reshape(-1, 5) # guard against no boxes
target = RBoxList(boxes, img.size, mode="xywha")
# classes
classes = [1] * len(anno)
classes = torch.tensor(classes)
target.add_field("labels", classes)
target.add_field("difficult", torch.tensor([0 for i in range(len(classes))]))
# masks
masks = [obj["points"].reshape((1, -1)).tolist() for obj in anno]
masks = SegmentationMask(masks, img.size)
target.add_field("masks", masks)
# target = target.clip_to_image(remove_empty=True)
if self.transforms is not None:
img, target = self.transforms(img, target)
assert target is not None, "{} target is None.".format(img_name)
return img, target, idx
def __getitem__(self, idx):
img = self.images[idx]
coco_idx = self.ids[idx]
anno = self.coco.loadAnns(self.coco.getAnnIds([coco_idx]))
boxes = [obj["bbox"] for obj in anno]
boxes = torch.as_tensor(boxes).reshape(-1, 4)
target = BoxList(boxes, img.size, mode="xywh").convert("xyxy")
classes = [obj["category_id"] for obj in anno]
classes = [
self.class_mapping[self.coco.cats[c]['name']] 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)
target.add_field("masks", masks)
target = target.crop([self.min_x, self.min_y, self.max_x,
self.max_y]).clip_to_image(remove_empty=True)
if self.transforms:
img, target = self.transforms(img, target)
return img, target, idx
def __getitem__(self, idx):
annotations = [
ann for ann in self.annt_labels['annotations']
if ann['image_id'] == idx
]
imageName = self.annt_labels["images"][idx]["file_name"]
bboxes, masks, classes = self._getannotationfields(annotations)
# opening the image
img = Image.open(imageName).convert("RGB")
bboxes = torch.as_tensor(bboxes).reshape(-1, 4)
target = BoxList(bboxes, img.size, mode='xywh').convert('xyxy')
classes = [self.json_category_id_to_contiguous_id[c] for c in classes]
classes = torch.tensor(classes)
target.add_field("labels", classes)
masks = SegmentationMask(masks, img.size, mode='poly')
target.add_field("masks", masks)
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 __getitem__(self, idx):
'''Load image.
Args:
idx: (int) image index.
Returns:
img: (tensor) image tensor.
loc_targets: (tensor) location targets.
cls_targets: (tensor) class label targets.
'''
# Load image and boxes.
img_path = self.img_pathes[idx]
if img_path not in self._cached_images.keys():
self._cached_images[img_path] = self.get_image(img_path)
img = self._cached_images[img_path]
boxes = self.boxes[idx].clone()
masks = self.masks[idx]
target = BoxList(boxes, img.size, mode="xyxy")
labels = self.labels[idx]
target.add_field("labels", labels)
masks = [[m] for m in masks]
masks = SegmentationMask(masks, img.size)
target.add_field("masks", masks)
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 __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)
target.add_field("masks", masks)
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 __getitem__(self, idx):
name = self.id_to_img_map[idx]
img_name = self.name2img_name(name)
img_path = os.path.join(self.img_root, img_name)
img = Image.open(img_path)
# TODO might be better to add an extra field
annos = self.annos[name]['annos']
boxes = [anno['box'] for anno in annos]
boxes = torch.as_tensor(boxes).reshape(-1, 4) # guard against no boxes
target = BoxList(boxes, img.size, mode="xyxy")
labels = [anno['label'] for anno in annos]
labels = torch.tensor(labels)
if not (labels < 3).all():
print(name)
target.add_field("labels", labels)
masks = [anno['segmentation'] for anno in annos]
masks = SegmentationMask(masks, img.size, mode='mask')
target.add_field("masks", masks)
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 __getitem__(self, idx):
img, anno = super(COCODataset, self).__getitem__(idx)
# 过滤 crowd annotations
# TODO might be better to add an extra field
anno = [obj for obj in anno if obj["iscrowd"] == 0]
# 把过滤之后的annotations对应的box转换成BoxList对象, 注意这些box可能属于多个类别
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")
# 将boxes的labels作为属性添加到target对象当中
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)
target.add_field("masks", masks)
# TODO: 关键点代码
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 __getitem__(self,idx):
img, anno = super(BirdDataset, self).__getitem__(idx)
# filter crowd annotations (doesn't apply to bird Dataset)
anno = [obj for obj in anno if obj["iscrowd"] == 0]
# load the image as a PIL image
boxes = [obj["bbox"] for obj in anno]
boxes = torch.as_tensor(boxes).reshape(-1, 4) # in case of no boxes
target = BoxList(boxes, img.size, mode="xywh").convert("xyxy")
classes = [obj["category_id"] for obj in anno]
#print(classes)
# Don't think I need this bit (and it's not coded)
#classes = [self.json_category_id_to_contiguous_id[c] for c in classes]
classes = torch.tensor(classes)
target.add_field("labels", classes)
if anno and "segmentation" in anno[0]:
masks = [obj["segmentation"] for obj in anno]
masks = SegmentationMask(masks, img.size, mode='poly')
target.add_field("masks", masks)
target = target.clip_to_image(remove_empty=True)
if self.transforms is not None:
img, target = self.transforms(img, target)
#print('pulling img',idx)
#print(target)
return img, target, idx
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)
target.add_field("masks", masks)
# VG related genome stuff
attributes = -torch.ones(
(len(boxes), self.max_attributes_per_ins), dtype=torch.long)
for idx, obj in enumerate(anno):
if "attribute_ids" in obj:
for jdx, att in enumerate(obj["attribute_ids"]):
attributes[idx, jdx] = att
target.add_field("attributes", attributes)
target = target.clip_to_image(remove_empty=True)
if self.transforms is not None:
# if self.opencv_loading:
## TEST only, Mimic Detectron Behaviour
# im = np.array(img).astype(np.float32)
# im = im[:, :, ::-1]
# im -= self.transforms.transforms[-1].mean
# im_shape = im.shape
# im_size_min = np.min(im_shape[0:2])
# im_size_max = np.max(im_shape[0:2])
# im_scale = float(800) / float(im_size_min)
# # Prevent the biggest axis from being more than max_size
# if np.round(im_scale * im_size_max) > 1333:
# im_scale = float(1333) / float(im_size_max)
# im = cv2.resize(
# im,
# None,
# None,
# fx=im_scale,
# fy=im_scale,
# interpolation=cv2.INTER_LINEAR
# )
# img = torch.from_numpy(im).permute(2, 0, 1)
# target.add_field("im_scale", im_scale)
img, target = self.transforms(img, target)
# else:
# img, target = self.transforms(img, target)
return img, target, idx
def _cat_mask(masks):
polygons_cat = []
size = masks[0].size
for mask in masks:
polygons = mask.get_polygons()
polygons_cat.extend(polygons)
masks_cat = SegmentationMask(polygons_cat, size)
return masks_cat
def __call__(self, image, target=None):
if target == None:
return image
if not self.do:
return image, target
w, h = image.size
if w != h:
return image, target
assert w == h
cx = w / 2
cy = cx
degree = random.uniform(0, 360)
radian = degree * math.pi / 180
new_image = image.rotate(-degree)
sin = math.sin(radian)
cos = math.cos(radian)
keypoints = target.get_field("keypoints")
polygons = list(
map(lambda x: x.polygons[0], keypoints.instances.polygons))
polygons = torch.stack(polygons, 0).reshape((-1, 2)).t()
M = torch.Tensor([[cos, -sin], [sin, cos]])
b = torch.Tensor([[(1 - cos) * cx + cy * sin],
[(1 - cos) * cy - cx * sin]])
new_points = M.mm(polygons) + b
new_points = new_points.t().reshape((-1, 8))
xmins, _ = torch.min(new_points[:, ::2], 1)
minx_idx = xmins < 1
xmins[minx_idx] = 1
ymins, _ = torch.min(new_points[:, 1::2], 1)
miny_idx = ymins < 1
ymins[miny_idx] = 1
xmaxs, _ = torch.max(new_points[:, ::2], 1)
xmax_idx = xmaxs > 1024
xmaxs[xmax_idx] = 1024
ymaxs, _ = torch.max(new_points[:, 1::2], 1)
ymax_idx = ymaxs > 1024
ymaxs[ymax_idx] = 1024
boxes = torch.stack([xmins, ymins, xmaxs, ymaxs], 1).reshape((-1, 4))
new_target = BoxList(boxes, image.size, mode="xyxy")
new_target._copy_extra_fields(target)
new_keypoints = SegmentationMask(new_points.reshape(
(-1, 1, 8)).tolist(),
image.size,
mode='poly')
new_target.add_field("keypoints", new_keypoints)
return new_image, new_target
def get_groundtruth(self, filename, width, height, aug_det):
anno = self._preprocess_annotation(self.ann_info[filename], aug_det)
target = BoxList(anno["boxes"], (width, height), mode="xywh").convert("xyxy")
target.add_field("labels", anno["labels"])
masks = SegmentationMask(anno["masks"], (width, height), type=self.mask_type)
target.add_field("masks", masks)
return target
def forward(self, anchors, box_cls, box_regression, coeffs, prototypes):
sampled_boxes = []
num_levels = len(box_cls)
anchors = list(zip(*anchors))
for a, c, r, co in zip(anchors, box_cls, box_regression, coeffs):
sampled_boxes.append(self.forward_for_single_feature_map(a, c, r, co))
boxlists = list(zip(*sampled_boxes))
boxlists = [cat_boxlist(boxlist) for boxlist in boxlists]
if num_levels > 1:
boxlists = self.select_over_all_levels(boxlists)
results = []
for prototypes_per_image, boxlists_per_image in zip(prototypes, boxlists):
coeffs_per_image = boxlists_per_image.get_field("coeffs")
# if DEBUG:
# print('range of prototypes_per_image:',\
# prototypes_per_image.min(), prototypes_per_image.max())
# assemble mask
masks_pred_per_image = prototypes_per_image.permute(1, 2, 0) @ coeffs_per_image.t()
masks_pred_per_image = masks_pred_per_image.permute(2, 0, 1)
masks_pred_per_image = self.mask_activation(masks_pred_per_image)
# crop
mask_h, mask_w = masks_pred_per_image.shape[1:]
resized_pred_bbox = boxlists_per_image.resize((mask_w, mask_h))
masks_pred_per_image = crop_zero_out(masks_pred_per_image, resized_pred_bbox.bbox)
# binarize
masks_pred_per_image = masks_pred_per_image > self.mask_threshold
# convert mask predictions to polygon format to save memory
if cfg.MODEL.YOLACT.CONVERT_MASK_TO_POLY:
cpu_device = torch.device("cpu")
masks_pred_per_image = SegmentationMask(masks_pred_per_image.to(cpu_device), \
(mask_w, mask_h), "mask")
if DEBUG:
print(len(masks_pred_per_image), mask_w, mask_h)
masks_pred_per_image = masks_pred_per_image.convert("poly")
else:
masks_pred_per_image = SegmentationMask(masks_pred_per_image, (mask_w, mask_h), "mask")
if DEBUG:
print(len(masks_pred_per_image), mask_w, mask_h)
# resize
img_w, img_h = boxlists_per_image.size
masks_pred_per_image = masks_pred_per_image.resize((img_w, img_h))
boxlists_per_image.add_field("masks", masks_pred_per_image)
results.append(boxlists_per_image)
return results
def __getitem__(self, index: int):
def folder(name):
return osj(self.root, f"{self.split}_{name}")
fname = self._files[index]
anno = self.anno[self.indice[index]]
image = pil_read(osj(folder("rgb"), fname))
objects = imread(osj(folder("object"), fname.replace("nyu_rgb_", "")))
n_obj = objects.max()
masks = [torch.from_numpy((objects == i).astype("int"))
for i in range(1, n_obj + 1)]
# in default image is concatenated into image
### TEST: no depth, no transform
#image_t = torch.from_numpy(image).permute(3, 0, 1)
boxes = torch.as_tensor(anno["bbox"]).reshape(-1, 4)
target = BoxList(boxes, image.size, mode="xyxy")
target.add_field("labels", torch.tensor(anno["label"]))
target.add_field("masks", SegmentationMask(masks, image.size, "mask"))
if self.use_depth:
if np.random.rand() < self.depth_prob:
depth = torch.zeros(objects.shape[0], objects.shape[1]).float()
else:
depth = torch.from_numpy(self.depth_data[index])
image_, target_, depth_ = self.transforms(image, target, depth)
image_ = torch.cat([image_, depth_], 0)
else:
image_, target_ = self.transforms(image, target)
if self.debug:
idx = self.split_indice[self.split][index]
print(self.indice[index], fname)
print("Image %s %f %f" % (str(image_.shape), image_[:3].min(), image_[:3].max()))
if self.use_depth:
print("Depth %f %f" % (image_[3].min(), image_[3].max()))
print("Mask", masks[0].shape, masks[0].min(), masks[0].max())
image_np = np.asarray(image).copy()
imwrite("%04d.png" % idx, image_np)
for i in range(n_obj):
box = anno["bbox"][i] # xmin, ymin, xmax, ymax
mask = masks[i].unsqueeze(2).numpy()
image_np[box[1],box[0]:box[2],:] = 255,0,0
image_np[box[3],box[0]:box[2],:] = 255,0,0
image_np[box[1]:box[3],box[0],:] = 255,0,0
image_np[box[1]:box[3],box[2],:] = 255,0,0
box_image = (image_np * mask)[box[1]:box[3], box[0]:box[2], :]
print(box, box_image.shape, image_np.shape, mask.shape)
print(box_image.max(), mask.max())
imwrite("%04d_%02d.png" % (idx, i), box_image)
imwrite("%04d_%02d_.png" % (idx, i), image_np * mask)
imwrite("%04d_.png" % idx, image_np)
return image_, target_, index
def forward_for_single_feature_map(self, pred, image_shapes):
"""
Arguments:
pred: tensor of size N, 1, H, W
"""
device = pred.device
# torch.cuda.synchronize()
# start_time = time.time()
bitmap = self.binarize(pred)
# torch.cuda.synchronize()
# end_time = time.time()
# print('binarize time:', end_time - start_time)
N, height, width = pred.shape[0], pred.shape[2], pred.shape[3]
# torch.cuda.synchronize()
# start_time = time.time()
bitmap_numpy = bitmap.cpu().numpy() # The first channel
pred_map_numpy = pred.cpu().numpy()
# torch.cuda.synchronize()
# end_time = time.time()
# print('gpu2numpy time:', end_time - start_time)
boxes_batch = []
rotated_boxes_batch = []
polygons_batch = []
scores_batch = []
# torch.cuda.synchronize()
# start_time = time.time()
for batch_index in range(N):
image_shape = image_shapes[batch_index]
boxes, scores, rotated_boxes, polygons = self.boxes_from_bitmap(
pred_map_numpy[batch_index],
bitmap_numpy[batch_index], width, height)
boxes = boxes.to(device)
if self.training and self.cfg.MODEL.SEG.AUG_PROPOSALS:
boxes = self.aug_tensor_proposals(boxes)
if boxes.shape[0] > self.top_n:
boxes = boxes[:self.top_n, :]
# _, top_index = scores.topk(self.top_n, 0, sorted=False)
# boxes = boxes[top_index, :]
# scores = scores[top_index]
# boxlist = BoxList(boxes, (width, height), mode="xyxy")
boxlist = BoxList(boxes, (image_shape[1], image_shape[0]), mode="xyxy")
if self.cfg.MODEL.SEG.USE_SEG_POLY or self.cfg.MODEL.ROI_BOX_HEAD.USE_MASKED_FEATURE or self.cfg.MODEL.ROI_MASK_HEAD.USE_MASKED_FEATURE:
masks = SegmentationMask(polygons, (image_shape[1], image_shape[0]))
boxlist.add_field('masks', masks)
boxlist = boxlist.clip_to_image(remove_empty=False)
# boxlist = remove_small_boxes(boxlist, self.min_size)
boxes_batch.append(boxlist)
rotated_boxes_batch.append(rotated_boxes)
polygons_batch.append(polygons)
scores_batch.append(scores)
# torch.cuda.synchronize()
# end_time = time.time()
# print('loop time:', end_time - start_time)
return boxes_batch, rotated_boxes_batch, polygons_batch, scores_batch
def get_item_for_softmax(self, idx):
iteration, idx, max_iter = idx['iteration'], idx['idx'], idx['max_iter']
if self.shuffle:
idx = self.shuffle[idx]
img, anno = self.get_image_ann(idx)
# we randomly select the max_box results. in the future, we should put it
# in the Transform
max_box = 800
# it will occupy 10G if it is 300 for one image.
if len(anno) > max_box:
logging.info('maximum box exceeds {} and we will truncate'.format(max_box))
import random
random.shuffle(anno)
anno = anno[:max_box]
anno = [o for o in anno if 'rect' in o]
if any('location_id' in a for a in anno):
# in this case, all locations should be unique for now.
assert len(set(a['location_id'] for a in anno)) == len(anno)
# coco data has this kind of property
anno = [obj for obj in anno if obj.get("iscrowd", 0) == 0]
anno = [a for a in anno if a['class'] in self.label_to_idx]
boxes = [obj["rect"] for obj in anno]
boxes = torch.as_tensor(boxes).reshape(-1, 4) # guard against no boxes
target = BoxList(boxes, img.size, mode="xyxy")
# 0 is the background
classes = [self.label_to_idx[obj["class"]] + 1 for obj in anno]
classes = torch.tensor(classes)
target.add_field("labels", classes)
self.add_tightness(target, anno)
if self.use_seg:
masks = [obj["segmentation"] for obj in anno]
from maskrcnn_benchmark.structures.segmentation_mask import SegmentationMask
masks = SegmentationMask(masks, img.size)
target.add_field("masks", masks)
target = target.clip_to_image(remove_empty=True)
if self.transforms is not None:
trans_input = {'image': img,
'rects': target,
'iteration': iteration,
'max_iter': max_iter,
'dataset': self,
}
trans_out = self.transforms(trans_input)
img, target = trans_out['image'], trans_out['rects']
return img, target, idx
def get_groundtruth(self, idx):
bboxes = []
masks = []
svg_name = self.data_dir / "svg" / "{:s}.svg".format(self.ids[idx])
with svg_name.open('r') as f_svg:
svg = f_svg.read()
global shift_x
global shift_y
global max_size
shift_max = int(self.canvas_size * 0.1)
shift_x = random.randint(-shift_max, shift_max)
shift_y = random.randint(-shift_max, shift_max)
max_size = self.canvas_size
def _replace_num(match):
x, y = [int(t) for t in match.group().split(" ")]
x = max(0, min(max_size - 1, x + shift_x))
y = max(0, min(max_size - 1, y + shift_y))
# y = max(900 - max_size, min(900, y + shift_y))
return "{:d} {:d}".format(x, y)
def _replace_path(match):
return re.sub(r'-?[0-9]+ [0-9]+', _replace_num, match.group())
# svg = re.sub(r'<path d=".*"></path>', _replace_path, svg)
# print(self.ids[idx], svg, "\n")
num_paths = len(et.fromstring(svg)[0])
for i in range(num_paths):
svg_xml = et.fromstring(svg)
svg_xml[0][0] = svg_xml[0][i]
del svg_xml[0][1:]
svg_one = et.tostring(svg_xml, method='xml')
# leave only one path
y_png = cairosvg.svg2png(bytestring=svg_one)
y_img = Image.open(io.BytesIO(y_png))
mask = (np.array(y_img)[:, :, 3] > 0)
try:
bboxes.append(self.get_bbox(mask, "xyxy"))
except IndexError:
continue
masks.append(mask.astype(np.uint8))
image_size_dict = self.get_img_info(idx)
image_size = (image_size_dict["height"], image_size_dict["width"])
boxlist = BoxList(bboxes, image_size, mode="xyxy")
boxlist.add_field("labels", torch.tensor([1] * len(bboxes)))
boxlist.add_field("mask", SegmentationMask(masks, image_size, "mask"))
return self.svg_to_png(svg), boxlist
