Update app.py

app.py

@@ -31,33 +31,59 @@ label_dict = {
     17: "Scarf",
 }
 
+# Function to process the image and generate the segmentation map
+
+
 # Function to process the image and generate the segmentation map
 def segment_image(image):
     # Perform segmentation
     result = pipe(image)
-    
+
     # Initialize an empty array for the segmentation map
     image_width, image_height = result[0]["mask"].size
     segmentation_map = np.zeros((image_height, image_width), dtype=np.uint8)
-    
+
     # Combine masks into a single segmentation map
-    for idx, entry in enumerate(result):
-        mask = np.array(entry["mask"])  # Convert the PIL mask to a NumPy array
-        segmentation_map[mask > 0] = idx  # Assign the class index
+    for entry in result:
+      label = entry["label"]  # Get the label (e.g., "Hair", "Upper-clothes")
+      mask = np.array(entry["mask"])  # Convert PIL Image to NumPy array
+      
+      # Find the index of the label in the original label dictionary
+      class_idx = [key for key, value in label_dict.items() if value == label][0]
+      
+      # Assign the correct class index to the segmentation map
+      segmentation_map[mask > 0] = class_idx
+    
     
+    # Get the unique class indices in the segmentation map
+    unique_classes = np.unique(segmentation_map)
+    # Print the names of the detected classes
+    print("Detected Classes:")
+    for class_idx in unique_classes:
+        print(f"- {label_dict[class_idx]}")
+
     # Create a matplotlib figure and visualize the segmentation map
     plt.figure(figsize=(8, 8))
     plt.imshow(segmentation_map, cmap="tab20")  # Visualize using a colormap
-    cbar = plt.colorbar(ticks=range(len(label_dict)), label="Classes")
-    cbar.ax.set_yticklabels([label_dict[i] for i in range(len(label_dict))])
-    plt.title("Combined Segmentation Map")
+    # Get the unique class indices in the segmentation map
+    unique_classes = np.unique(segmentation_map)
+
+
+    # Filter the label dictionary to include only detected classes
+    filtered_labels = {idx: label_dict[idx] for idx in unique_classes}
+
+    # Create a dynamic colorbar with only the detected classes
+    cbar = plt.colorbar(ticks=unique_classes)
+    cbar.ax.set_yticklabels([filtered_labels[i] for i in unique_classes])
+    plt.title("Segmented Image with Detected Classes")
     plt.axis("off")
-    
-    # Save the figure as a PIL image for Gradio
-    plt.savefig("segmented_output.png", bbox_inches="tight")  # Save as a temporary file
-    plt.close()  # Close the figure to free memory
+    plt.savefig("segmented_output.png", bbox_inches="tight")
+    plt.close()
     return Image.open("segmented_output.png")
 
+
+
+
 # Gradio interface
 interface = gr.Interface(
     fn=segment_image,