WaterKnight1998 / SemTorch

Apache License 2.0
162 stars 15 forks source link

Instance segmentation: support for multiple masks in the same image #8

Open ZenBel opened 3 years ago

ZenBel commented 3 years ago

Hello,

I am trying to use SemTorch for instance segmentation for the case when multiple masks are present in the same image. By looking at the source code and the MaskRCNN notebook , it seems that only one mask per image is supported. Am I correct?

I have written a small function that builds the bounding boxes for all the masks in an image, and assigns the corresponding (binary) labels:

def get_bboxes(o):

    # Read image and corresponding mask
    img_0 = cv2.imread(str(o))
    mask_0 = cv2.imread(str(get_msk(o)))
    # Change to grayscale for finding contours
    mask_gray = cv2.cvtColor(mask_0, cv2.COLOR_BGR2GRAY)
    ret, thresh = cv2.threshold(mask_gray, 0, 255, cv2.THRESH_BINARY)

    # Find contours
    contours, hierarchy = cv2.findContours(image=thresh, mode=cv2.RETR_EXTERNAL, method=cv2.CHAIN_APPROX_NONE)

    bboxes = []
    for cont in contours:
      xmin = cont[:,:,0].min()
      ymin = cont[:,:,1].min()
      xmax = cont[:,:,0].max()
      ymax = cont[:,:,1].max()
      bboxes.append([xmin, ymin, xmax, ymax])

    cat = [1]*len(bboxes)

    return TensorBBox.create(bboxes), TensorCategory(cat)

but when I run

def get_dict(o):
    return {"boxes": get_bboxes(o)[0], "labels": get_bboxes(o)[1], "masks": get_msk(o)}

getters = [lambda o: o, get_dict]

maskrccnnDataBlock = DataBlock(
    blocks=(ImageBlock, MaskRCNNBlock),
    get_items=get_image_files,
    getters=getters,
    splitter=RandomSplitter(valid_pct=0.2, seed=42),
    item_tfms=[IntToFloatTensorMaskRCNN],
    dl_type=TfmdDLV2,
    n_inp=1
)

maskrccnnDataBlock.summary(path_im)

I get the following error from .summary():

Collating items in a batch
Error! It's not possible to collate your items in a batch
Could not collate the 0-th members of your tuples because got the following shapes
torch.Size([3, 305, 305]),torch.Size([3, 305, 305]),torch.Size([3, 305, 305]),torch.Size([3, 305, 305])
---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
<ipython-input-137-c9f5d49c6974> in <module>()
----> 1 maskrccnnDataBlock.summary(path_im)
      2 # print("Batch Size {}".format(bs))
[...]
/usr/local/lib/python3.7/dist-packages/torch/_tensor.py in __torch_function__(cls, func, types, args, kwargs)
   1021 
   1022         with _C.DisableTorchFunction():
-> 1023             ret = func(*args, **kwargs)
   1024             return _convert(ret, cls)
   1025 

RuntimeError: stack expects each tensor to be equal size, but got [4, 4] at entry 0 and [2, 4] at entry 1

which is probably due to the fact that in one image there are 4 masks, and in the other only 2. Any idea on how to go about this issue?

Thanks,

Zeno

cdhasselerharm commented 2 years ago

Hey @ZenBel Did you manage to overcome this issue? Interested in doing the same thing

AhmadZobairSurosh commented 2 years ago

@ZenBel @Chris-Has @WaterKnight1998

HI!

did you manage to do muli mask per image segmentation for Maskrcnn?

Thanks in advance