mattmdjaga

/

segformer_b2_clothes

like 356

Hi, so to get a dictionary of the ids which are used in the output to labels you can do "model.config.id2label" and that will give you a mapping of the prediction labels to string labels. I'm not sure what you mean getting the sizes of the labels though, could you elaborate?

@mattmdjaga how do I get the labels corresponding to the segment colors? After running the code in the model card

I'm not sure how to get the colour map from the example code but gpt 4 helped me get this which sets a new colourmap and lets you get a dictionary of label to colour

import matplotlib as mpl
label_names = list(model.config.id2label)
# Create a color map with the same number of colors as your labels
# Use the updated method to get the colormap
cmap = mpl.colormaps['tab20']

# Create the figure and axes for the plot and the colorbar
fig, ax = plt.subplots()

# Display the segmentation
im = ax.imshow(pred_seg, cmap=cmap)

# Create a colorbar
cbar = fig.colorbar(im, ax=ax, ticks=range(len(label_names)))
cbar.ax.set_yticklabels(label_names)

plt.show()

# Get the number of labels
n_labels = len(label_names)

# Extract RGB values for each color in the colormap
colors = cmap.colors[:n_labels]

# Convert RGBA to RGB by omitting the Alpha value
rgb_colors = [color[:3] for color in colors]

# Create a dictionary mapping labels to RGB colors
label_to_color = dict(zip(label_names, rgb_colors))

# Display the mapping
for label, color in label_to_color.items():
    print(f"{label}: {color}")

Thanks!

Hi! Thank you for the help! To answer your question and clarify what I was saying: I'm essentially trying to get the upper left plot point and the lower right plot point of a layer as to get the size of the clothing piece in the image. So for example if you had an image of a shirt I would want the size of the shirt in relation to the image, how would I go about doing that?

You can get the % of pixels predicted for each label using the code below:

id2label = model.config.id2label
labels = pred_seg.unique()
total_pixels = pred_seg.shape[0] * pred_seg.shape[1]
for label in labels:
    mask_val = (pred_seg==label).sum()
    print(f"{id2label[label.item()]} takes {(mask_val/total_pixels)*100:.3f}%")

For the model card example that should return:

Background takes 77.867%
Hair takes 0.352%
Upper-clothes takes 13.596%
Pants takes 6.270%
Face takes 1.519%
Right-arm takes 0.221%
Bag takes 0.175%

Hope this helps