Spaces:
Runtime error
Runtime error
File size: 3,459 Bytes
688c60a 394311c 688c60a 394311c f16d728 688c60a 394311c 688c60a 79d5975 688c60a 394311c 688c60a |
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 |
import torch
from torchvision.models.detection import fasterrcnn_resnet50_fpn, FastRCNNPredictor
from torchvision.transforms import ToTensor
from PIL import Image, ImageDraw, ImageFont
import gradio as gr
import io
import pandas as pd
# Define classes (COCO-style)
COCO_CLASSES = [
'__background__', 'person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus',
'train', 'truck', 'boat', 'traffic light', 'fire hydrant', 'N/A', 'stop sign',
'parking meter', 'bench', 'bird', 'cat', 'dog', 'horse', 'sheep', 'cow',
'elephant', 'bear', 'zebra', 'giraffe', 'N/A', 'backpack', 'umbrella', 'N/A',
'N/A', 'handbag', 'tie', 'suitcase', 'frisbee', 'skis', 'snowboard',
'sports ball', 'kite', 'baseball bat', 'baseball glove', 'skateboard',
'surfboard', 'tennis racket', 'bottle', 'N/A', 'wine glass', 'cup', 'fork',
'knife', 'spoon', 'bowl', 'banana', 'apple', 'sandwich', 'orange',
'broccoli', 'carrot', 'hot dog', 'pizza', 'donut', 'cake', 'chair', 'couch',
'potted plant', 'bed', 'N/A', 'dining table', 'N/A', 'N/A', 'toilet', 'N/A',
'tv', 'laptop', 'mouse', 'remote', 'keyboard', 'cell phone', 'microwave',
'oven', 'toaster', 'sink', 'refrigerator', 'N/A', 'book', 'clock', 'vase',
'scissors', 'teddy bear', 'hair drier', 'toothbrush'
]
# Load model
def load_model(checkpoint_path=None, num_classes=91):
model = fasterrcnn_resnet50_fpn(pretrained=True)
if checkpoint_path:
# Replace head for custom classes if needed
in_features = model.roi_heads.box_predictor.cls_score.in_features
model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)
model.load_state_dict(torch.load(checkpoint_path, map_location="cpu"))
model.eval()
return model
model = load_model("frcnn_model.pth") # Replace with None to use default COCO model
# Run inference
def detect_objects(image):
image = image.convert("RGB")
image_tensor = ToTensor()(image).unsqueeze(0)
with torch.no_grad():
outputs = model(image_tensor)
boxes = outputs[0]['boxes']
labels = outputs[0]['labels']
scores = outputs[0]['scores']
threshold = 0.5
draw = ImageDraw.Draw(image)
font = ImageFont.load_default()
log_data = []
for box, label, score in zip(boxes, labels, scores):
if score >= threshold:
box = box.tolist()
label_name = COCO_CLASSES[label.item()]
draw.rectangle(box, outline="red", width=2)
draw.text((box[0], box[1]), f"{label_name} ({score:.2f})", fill="yellow", font=font)
log_data.append({"Object": label_name, "Score": round(score.item(), 2)})
# Prepare image for download
buffer = io.BytesIO()
image.save(buffer, format="PNG")
buffer.seek(0)
# Prepare table
table = pd.DataFrame(log_data) if log_data else pd.DataFrame(columns=["Object", "Score"])
return image, buffer, table
# Gradio interface
app = gr.Interface(
fn=detect_objects,
inputs=gr.Image(type="pil", label="Upload Image"),
outputs=[
gr.Image(type="pil", label="Detected Image"),
gr.File(label="Download Image"),
gr.Dataframe(headers=["Object", "Score"], label="Detection Log")
],
title="🧠 Object Detection App (Faster R-CNN)",
description="Upload an image to detect objects using a pretrained or custom Faster R-CNN model. View logs and download the result."
)
if __name__ == "__main__":
app.launch() |