sandbox338/objectdetection

1import streamlit as st2from PIL import Image3import numpy as np4import torch5import asyncio6from detectron2.engine import DefaultPredictor7from detectron2.config import get_cfg8from detectron2 import model_zoo9from detectron2.utils.visualizer import Visualizer10from detectron2.data import MetadataCatalog11 12# Fix event loop issue13try:14    asyncio.get_running_loop()15except RuntimeError:16    asyncio.set_event_loop(asyncio.new_event_loop())17 18# Title and uploader19st.title("Detectron2 Object Detection")20st.write("Upload an image to perform object detection")21 22uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])23 24@st.cache_resource25def load_model():26    cfg = get_cfg()27    cfg.merge_from_file(model_zoo.get_config_file("COCO-Detection/faster_rcnn_R_50_FPN_3x.yaml"))28    cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.529    cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url("COCO-Detection/faster_rcnn_R_50_FPN_3x.yaml")30    predictor = DefaultPredictor(cfg)31    return predictor32 33def predict_fn(predictor, image):34    image_array = np.array(image)[:, :, :3]35    outputs = predictor(image_array)36    return outputs["instances"], image_array37 38def visualize_predictions(image, instances):39    v = Visualizer(image[:, :, ::-1], MetadataCatalog.get("coco_2017_val"), scale=1.2)40    v = v.draw_instance_predictions(instances)41    result = v.get_image()42    return result[:, :, ::-1]43 44if uploaded_file is not None:45    image = Image.open(uploaded_file)46    st.image(image, caption="Uploaded Image", use_column_width=True)47    st.write("Processing...")48 49    predictor = load_model()50    instances, image_array = predict_fn(predictor, image)51    result_image = visualize_predictions(image_array, instances)52 53    st.image(result_image, caption="Detected Objects", use_column_width=True)54