Tarun77/ObjectDetection

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