Create app.py

app.py

@@ -0,0 +1,93 @@
+import torch
+from torchvision import transforms
+from PIL import Image, ImageDraw, ImageEnhance
+import requests
+from torchvision.models.detection import maskrcnn_resnet50_fpn
+import random
+
+# Load the Mask R-CNN model
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+model = maskrcnn_resnet50_fpn(pretrained=True).to(device).eval()
+
+# Function to preprocess the image
+def preprocess_image(image_path):
+    # Open and convert to RGB
+    image = Image.open(image_path).convert("RGB")  
+    transform = transforms.Compose([
+        # Convert image to a tensor
+        transforms.ToTensor(),  
+    ])
+    # Add batch dimension and send to device
+    return transform(image).unsqueeze(0).to(device), image  
+
+# Run object detection
+def detect_objects(image_path, threshold=0.5):
+    image_tensor, image_pil = preprocess_image(image_path)
+    with torch.no_grad():
+        outputs = model(image_tensor)[0]  # Get model output
+    
+    # Extract data from model output
+    masks = outputs["masks"]  # Object masks
+    labels = outputs["labels"]  # Object labels
+    scores = outputs["scores"]  # Confidence scores
+    filtered_masks = []
+    
+    for i in range(len(masks)):
+        # Only keep objects with high confidence
+        if scores[i] >= threshold:  
+            # Convert to binary mask
+            mask = masks[i, 0].mul(255).byte().cpu().numpy()  
+            filtered_masks.append((mask, labels[i].item(), scores[i].item()))
+    
+    return filtered_masks, image_pil
+
+# Apply color masks to detected objects
+def apply_instance_masks(image_path):
+    masks, image = detect_objects(image_path)
+
+    # Convert to RGBA to support transparency
+    img = image.convert("RGBA")  
+    # Create a transparent layer
+    overlay = Image.new("RGBA", img.size, (0, 0, 0, 0))  
+    draw = ImageDraw.Draw(overlay)
+
+    # Store unique colors for each object category
+    color_map = {}  
+
+    for mask, label, score in masks:
+        if label not in color_map:
+            # Assign a random color for this object category
+            color_map[label] = (random.randint(50, 50), random.randint(225, 255), random.randint(50, 50), 150)
+
+        mask_pil = Image.fromarray(mask, mode="L")  # Convert mask to grayscale image
+        colored_mask = Image.new("RGBA", mask_pil.size, color_map[label])  # Create a color mask
+        overlay.paste(colored_mask, (0, 0), mask_pil)  # Apply mask to the overlay
+    
+    # Combine the original image with the overlay
+    result_image = Image.alpha_composite(img, overlay)
+
+    return result_image.convert("RGB")  # Convert back to RGB mode
+
+import gradio as gr
+
+with gr.Blocks() as demo:
+    gr.Markdown("## Object Detection with Mask R-CNN")
+    gr.Markdown("This demo applies instance segmentation to an image using Mask R-CNN.")
+    
+    with gr.Row():
+        with gr.Column():
+            input_image = gr.Image(label="Input Image", type="filepath")
+            threshold = gr.Slider(minimum=0.0, maximum=1.0, value=0.5, label="Confidence Threshold")
+            detect_button = gr.Button("Detect Objects")
+        
+        with gr.Column():
+            output_image = gr.Image(label="Output Image with Masks")
+
+    detect_button.click(
+        fn=lambda img_path, thresh: apply_instance_masks(img_path) if img_path else None,
+        inputs=[input_image, threshold],
+        outputs=output_image
+    )
+
+demo.launch()