File size: 2,710 Bytes
96f5c80 5d6db6d 96f5c80 504b1c1 5d6db6d 504b1c1 96f5c80 504b1c1 5d6db6d 96f5c80 504b1c1 7b90a74 504b1c1 7b90a74 504b1c1 7b90a74 729d3e8 504b1c1 729d3e8 504b1c1 96f5c80 504b1c1 5d6db6d 729d3e8 504b1c1 7b90a74 504b1c1 7b90a74 504b1c1 5d6db6d 504b1c1 5d6db6d 729d3e8 |
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 |
import streamlit as st
from PIL import Image
import torch
import torchvision.transforms as T
import torchvision
import matplotlib.pyplot as plt
import matplotlib.patches as patches
# Load pre-trained object detection model
model = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=True)
model.eval()
# COCO class labels
COCO_INSTANCE_CATEGORY_NAMES = [
'__background__', 'person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus',
'train', 'truck', 'boat', 'traffic light', 'fire hydrant', 'stop sign',
'parking meter', 'bench', 'bird', 'cat', 'dog', 'horse', 'sheep', 'cow',
'elephant', 'bear', 'zebra', 'giraffe', 'backpack', 'umbrella', 'handbag',
'tie', 'suitcase', 'frisbee', 'skis', 'snowboard', 'sports ball', 'kite',
'baseball bat', 'baseball glove', 'skateboard', 'surfboard', 'tennis racket',
'bottle', 'wine glass', 'cup', 'fork', 'knife', 'spoon', 'bowl', 'banana',
'apple', 'sandwich', 'orange', 'broccoli', 'carrot', 'hot dog', 'pizza',
'donut', 'cake', 'chair', 'couch', 'potted plant', 'bed', 'dining table',
'toilet', 'tv', 'laptop', 'mouse', 'remote', 'keyboard', 'cell phone',
'microwave', 'oven', 'toaster', 'sink', 'refrigerator', 'book', 'clock',
'vase', 'scissors', 'teddy bear', 'hair drier', 'toothbrush'
]
# Streamlit UI
st.set_page_config(page_title="AI Object Detector", layout="centered")
st.title("🎯 Object Detection with AI")
st.markdown("Upload an image and let the AI detect objects with names!")
uploaded_file = st.file_uploader("📷 Upload an image", type=["jpg", "png", "jpeg"])
if uploaded_file:
image = Image.open(uploaded_file).convert("RGB")
st.image(image, caption="Uploaded Image", use_column_width=True)
# Transform image for model input
transform = T.Compose([T.ToTensor()])
image_tensor = transform(image).unsqueeze(0)
with st.spinner("🔍 Detecting objects..."):
predictions = model(image_tensor)[0]
# Draw bounding boxes
fig, ax = plt.subplots(1)
ax.imshow(image)
threshold = 0.7 # confidence threshold
for idx in range(len(predictions["boxes"])):
score = predictions["scores"][idx].item()
if score > threshold:
box = predictions["boxes"][idx].detach().numpy()
label = COCO_INSTANCE_CATEGORY_NAMES[predictions["labels"][idx]]
x1, y1, x2, y2 = box
rect = patches.Rectangle((x1, y1), x2 - x1, y2 - y1,
linewidth=2, edgecolor='lime', facecolor='none')
ax.add_patch(rect)
ax.text(x1, y1 - 10, f"{label} ({score:.2f})", color='lime',
fontsize=10, backgroundcolor='black')
st.pyplot(fig)
|