mattmdjaga

/

segformer_b2_clothes

like 336

So the output is a 2D image where each pixel is the prediction class, and the classes are "0": "Background", "1": "Hat", "2": "Hair", "3": "Sunglasses", "4": "Upper-clothes", "5": "Skirt", "6": "Pants", "7": "Dress", "8": "Belt", "9": "Left-shoe", "10": "Right-shoe", "11": "Face", "12": "Left-leg", "13": "Right-leg", "14": "Left-arm", "15": "Right-arm", "16": "Bag", "17": "Scarf". So you can't segment "car" using this model, but you can change the mask to only have one class by doing something like output[output != 1] = 0, and this would change everything that isn't a hat to 0.
Hope that this answers your question.

Sorry for the naive question , the api gives
pred_seg = upsampled_logits.argmax(dim=1)[0]
plt.imshow(pred_seg)
image with all the segmented class as output,
what should i do to get only "SKIRT" mask ?

try this in the line before plt.imshow(pred_seg) :

pred_seg[pred_seg != 5] = 0

This creates a tensor of size [class heightwidth] . which i am unable to convert to PIL image. how this can be done

sorry for the delay, try:
arr_seg = pred_seg.cpu().numpy().astype("uint8")
pil_seg = Image.fromarray(arr_seg)

pil_seg should be a pil image, keep in mind that the values are between 0-17, so you might need to do something like
arr_seg *= 255 or arr_seg[arr_seg > 0] = 255
in order to see the mask in a pil image

Thanks @mattmdjaga that works