WaterKnight
(David Lacalle Castillo)
May 2, 2020, 2:20pm
15
My DataBlock is eassier, is just a dict.
I am using it just for the target and keeping ImageBlock for the input.
It works pretty straidforward. I am not using transforms like resize and normalize. It is done in first layer of torchvision.models.detection.maskrcnn_resnet50_fpn.
class MaskRCNN(dict): @classmethod def create(cls, dictionary): return cls(dict({x:dictionary for x in dictionary.keys()})) def show(self, ctx=None, **kwargs): dictionary = self boxes = dictionary[“boxes”] labels = dictionary[“labels”] masks = dictionary[“masks”] result = masks return show_image(result, ctx=ctx, **kwargs) def MaskRCNNBlock(): return TransformBlock(type_tfms=MaskRCNN.create, batch_tfms=IntToFloatTensor) def get_bbox(o): label_path = get_y_fn(o) mask=PILMask.create(label_path) pos = np.where(mask) xmin = np.min(pos[1]) xmax = np.max(pos[1]) ymin = np.min(pos[0]) ymax = np.max(pos[0]) return TensorBBox.create([xmin, ymin, xmax, ymax]) def get_bbox_label(o): return TensorCategory([1]) def get_mask(o): label_path = get_y_fn(o) mask=PILMask.create(label_path) mask=image2tensor(mask) return TensorMask(mask) def get_dict(o): return {“boxes”: get_bbox(o), “labels”: get_bbox_label(o),“masks”: get_mask(o)} getters = [lambda o: o, get_dict] maskrccnnDataBlock = DataBlock( blocks=(ImageBlock, MaskRCNNBlock), get_items=partial(get_image_files,folders=[manual_name]), getters=getters, splitter=RandomSplitter(valid_pct=0.1,seed=2020), item_tfms=Resize((size,size)), batch_tfms=Normalize.from_stats(*imagenet_stats) ) maskrccnnDataBlock.summary(path_images) dls = maskrccnnDataBlock.dataloaders(path_images,bs=bs)