mattmdjaga

/

segformer_b2_clothes

like 359

I'm not sure if this is what you mean but, here is code to take a prediction and save each of the detected items and save them as an image.

# use the started code provided in the model card up to the second last line where you get the variable "pred_seg"
import numpy as np 
import torch

segments = torch.unique(pred_seg) # Get a list of all the predicted items
for i in segments: 
    mask = pred_seg == i # Filter out anything that isn't the current item
    img = Image.fromarray((mask * 255).numpy().astype(np.uint8))
    name = model.config.id2label[i.item()] # get the item name
    img.save(f"{name}.png")

Let me know if anything doesn't work or this is what you wanted

I'm not sure if this is what you mean but, here is code to take a prediction and save each of the detected items and save them as an image.

# use the started code provided in the model card up to the second last line where you get the variable "pred_seg"
import numpy as np 
import torch

segments = torch.unique(pred_seg) # Get a list of all the predicted items
for i in segments: 
    mask = pred_seg == i # Filter out anything that isn't the current item
    img = Image.fromarray((mask * 255).numpy().astype(np.uint8))
    name = model.config.id2label[i.item()] # get the item name
    img.save(f"{name}.png")

Let me know if anything doesn't work or this is what you wanted

hi,
thanks alot for your reply, this is exactly what i wanted.
i really appreciate your work in this model.

Hi, sorry to reopen this discussion. The code so far has been really helpful but I was wondering if there is a way to show the original image with the masking. For example an image of a man wearing a shirt and trousers. How would I go about using a mask so that I could just show the shirt from the original image?

Hi, sorry to reopen this discussion. The code so far has been really helpful but I was wondering if there is a way to show the original image with the masking. For example an image of a man wearing a shirt and trousers. How would I go about using a mask so that I could just show the shirt from the original image?

+1

So the code snippet below uses the segmentation to turn anything that isn't 4 (upper clothes) into a black pixel. Is this what you were looking for?

import numpy as np
seg_arr = pred_seg.numpy()
seg_arr = seg_arr.astype(np.uint8)
img = np.array(image)
img[seg_arr!=4] = 0
plt.imshow(img)

That's exactly what I was looking for. Thank you!