Create app.py

e07a8a4 verified over 1 year ago

3.15 kB

import gradio as gr
from transformers import pipeline
from PIL import Image
import numpy as np
import matplotlib.pyplot as plt

# Load the segmentation pipeline
pipe = pipeline("image-segmentation", model="mattmdjaga/segformer_b2_clothes")



# Your predefined label dictionary
label_dict = {
    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",
}

# 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 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
    # 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")
    plt.savefig("segmented_output.png", bbox_inches="tight")
    plt.close()
    return Image.open("segmented_output.png")




# Gradio interface
interface = gr.Interface(
    fn=segment_image,
    inputs=gr.Image(type="pil"),  # Input is an image
    outputs=gr.Image(type="pil"),  # Output is an image with the colormap
    #examples=["example_image.jpg"],  # Use the saved image as an example
    examples=["1.jpg", "2.jpg", "3.jpg"],
    title="Clothes Segmentation with Colormap",
    description="Upload an image, and the segmentation model will produce an output with a colormap applied to the segmented classes."
)

# Launch the app
interface.launch()

	import gradio as gr
	from transformers import pipeline
	from PIL import Image
	import numpy as np
	import matplotlib.pyplot as plt

	# Load the segmentation pipeline
	pipe = pipeline("image-segmentation", model="mattmdjaga/segformer_b2_clothes")



	# Your predefined label dictionary
	label_dict = {
	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",
	}

	# 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 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
	# 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")
	plt.savefig("segmented_output.png", bbox_inches="tight")
	plt.close()
	return Image.open("segmented_output.png")




	# Gradio interface
	interface = gr.Interface(
	fn=segment_image,
	inputs=gr.Image(type="pil"), # Input is an image
	outputs=gr.Image(type="pil"), # Output is an image with the colormap
	#examples=["example_image.jpg"], # Use the saved image as an example
	examples=["1.jpg", "2.jpg", "3.jpg"],
	title="Clothes Segmentation with Colormap",
	description="Upload an image, and the segmentation model will produce an output with a colormap applied to the segmented classes."
	)

	# Launch the app
	interface.launch()