Spaces:
Sleeping
Sleeping
File size: 3,878 Bytes
7811fb5 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 |
import torch
import gradio as gr
import torchvision.transforms as T
from PIL import Image, ImageDraw, ImageOps
import numpy as np
from torchvision.models.detection import maskrcnn_resnet50_fpn
from torchvision.models.detection.faster_rcnn import FastRCNNPredictor
from torchvision.models.detection.mask_rcnn import MaskRCNNPredictor
import os
# Set up device
device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
# Load and configure the Mask R-CNN model with 2 classes
model_path = "mask_rcnn_lego.pth"
if not os.path.exists(model_path):
raise FileNotFoundError(
"The model file 'mask_rcnn_lego.pth' was not found in the directory."
)
model = maskrcnn_resnet50_fpn(weights="DEFAULT")
in_features = model.roi_heads.box_predictor.cls_score.in_features
# Update the box predictor head to match 2 classes (background + LEGO)
model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes=2)
# Update the mask predictor head to match 2 classes
in_features_mask = model.roi_heads.mask_predictor.conv5_mask.in_channels
hidden_layer = 256
model.roi_heads.mask_predictor = MaskRCNNPredictor(
in_features_mask, hidden_layer, num_classes=2
)
# Now, load the state_dict for your custom model
model.load_state_dict(torch.load(model_path, map_location=device))
model.to(device)
model.eval()
# Set up transformations
transform = T.Compose([T.ToTensor()])
# Function to create pseudo-masks based on bounding boxes
def create_pseudo_mask(image, box):
mask = Image.new("L", image.size, 0) # Create a blank mask
draw = ImageDraw.Draw(mask)
draw.rectangle(box, fill=255) # Fill in the bounding box area
return mask
# Function to process image with pseudo-mask visualization and bounding boxes
def detect_legos(image, use_pseudo_masks=True):
# Apply transformations
img_tensor = transform(image).unsqueeze(0).to(device)
# Make predictions with the custom model
with torch.no_grad():
outputs = model(img_tensor)
# Extract boxes and scores above threshold
boxes = outputs[0]["boxes"].cpu().numpy()
scores = outputs[0]["scores"].cpu().numpy()
thresholded_indices = [i for i, score in enumerate(scores) if score >= 0.5]
boxes = boxes[thresholded_indices]
num_legos_detected = len(boxes)
# Draw pseudo-masks on the image first
image_with_masks = image.copy()
for box in boxes:
x1, y1, x2, y2 = box
# Use pseudo-masks based on bounding boxes
if use_pseudo_masks:
mask_img = create_pseudo_mask(image, [x1, y1, x2, y2])
mask_img = ImageOps.colorize(
mask_img.convert("L"), black="blue", white="blue"
).convert("RGBA")
image_with_masks.paste(mask_img, (0, 0), mask_img)
# Draw the bounding boxes on top of the masks for better visibility
draw = ImageDraw.Draw(image_with_masks)
for box in boxes:
x1, y1, x2, y2 = box
draw.rectangle(
[x1, y1, x2, y2], outline="yellow", width=3
) # Draw yellow bounding box
# Set title with count of detected LEGO pieces
title = f"Detected LEGO pieces: {num_legos_detected}"
return image_with_masks, title
# Gradio interface function
def gradio_interface(image):
image_with_masks, title = detect_legos(image, use_pseudo_masks=True)
return image_with_masks, title
# Set up Gradio Interface
interface = gr.Interface(
fn=gradio_interface,
inputs=gr.Image(type="pil"),
outputs=[gr.Image(type="pil"), gr.Textbox(label="Detection Summary")],
title="LEGO Detection with Mask R-CNN",
description="Upload an image to detect and count LEGO pieces with bounding boxes and simulated masks.",
)
# Launch the Gradio app
interface.launch()
|