Tarun77/ObjectDetection

1import streamlit as st2from ultralytics import YOLO3from PIL import Image4import cv25import numpy as np6import pandas as pd7import tempfile8 9 10import os11os.environ["YOLO_CONFIG_DIR"] = "/tmp/ultralytics"12os.environ["XDG_CONFIG_HOME"] = "/tmp"13 14os.environ["STREAMLIT_HOME"] = "/tmp/.streamlit"15 16model = YOLO('/best.pt')17class_names = model.names  # or your own list18 19st.title("Custom Object Detection Demo")20 21media_type = st.radio("Choose media type:", ["Image", "Video"])22 23if media_type == "Image":24    uploaded_file = st.file_uploader("Upload an image", type=["jpg", "png"])25    26    if uploaded_file:27        image = Image.open(uploaded_file).convert("RGB")28        results = model.predict(image)29        plotted_img = results[0].plot()30        plotted_img = cv2.cvtColor(plotted_img, cv2.COLOR_BGR2RGB)31 32        # Show images side by side33        col1, col2 = st.columns(2)34        with col1:35            st.image(image, caption="📷 Original Image", use_column_width=True)36        with col2:37            st.image(plotted_img, caption="Detection Output", use_column_width=True)38 39        # Show detected objects in a table40        data = []41        for box in results[0].boxes:42            cls_id = int(box.cls[0])43            conf = float(box.conf[0])44            data.append({45                "Class": class_names[cls_id],46                "Confidence (%)": f"{conf * 100:.2f}"47            })48        if data:49            df = pd.DataFrame(data)50            df.index = df.index + 151            df.index.name = "S. No."52            st.subheader("Detected Objects")53            st.table(df)54        else:55            st.info("No objects detected.")56 57elif media_type == "Video":58    uploaded_video = st.file_uploader("Upload a video", type=["mp4", "mov", "avi"])59    60    if uploaded_video:61        tfile = tempfile.NamedTemporaryFile(delete=False)62        tfile.write(uploaded_video.read())63        64        cap = cv2.VideoCapture(tfile.name)65        stframe = st.empty()66 67        st.info("Processing video...")68 69        while cap.isOpened():70            ret, frame = cap.read()71            if not ret:72                break73 74            results = model.predict(frame)75            plotted_frame = results[0].plot()76            plotted_frame = cv2.cvtColor(plotted_frame, cv2.COLOR_BGR2RGB)77 78            stframe.image(plotted_frame, channels="RGB", use_column_width=True)79 80        cap.release()81        st.success("Video processing completed.")82