Create streamlitUI.py

streamlitUI.py

@@ -0,0 +1,155 @@
+import os
+import cv2
+import tempfile
+import requests
+import base64
+import numpy as np
+from PIL import Image
+from io import BytesIO
+from ultralytics import YOLO
+import streamlit as st
+import yt_dlp as youtube_dl
+
+def load_yolov8_model(model_name='yolov8s.pt'):
+    try:
+        return YOLO(model_name)
+    except Exception as e:
+        st.error(f"Error loading model: {e}")
+        return None
+
+def detect_objects(image, model):
+    try:
+        results = model(image)
+        for result in results[0].boxes:
+            x1, y1, x2, y2 = map(int, result.xyxy[0])
+            label = model.names[int(result.cls)]
+            confidence = result.conf.item()
+            cv2.rectangle(image, (x1, y1), (x2, y2), (0, 255, 0), 2)
+            label_text = f'{label} {confidence:.2f}'
+            cv2.putText(image, label_text, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)
+        return image
+    except Exception as e:
+        st.error(f"Error during object detection: {e}")
+        return None
+
+def process_image_with_yolov8(model_name, image=None, url=None):
+    model = load_yolov8_model(model_name)
+    if model is None:
+        return None
+
+    if url:
+        if url.startswith('data:image'):
+            try:
+                header, encoded = url.split(',', 1)
+                image_data = base64.b64decode(encoded)
+                image = Image.open(BytesIO(image_data))
+            except Exception as e:
+                st.error(f"Error decoding base64 image: {e}")
+                return None
+        else:
+            try:
+                response = requests.get(url)
+                response.raise_for_status()
+                image = Image.open(BytesIO(response.content))
+            except Exception as e:
+                st.error(f"Error loading image from URL: {e}")
+                return None
+
+    try:
+        image = np.array(image)
+        output_image = detect_objects(image, model)
+        return output_image
+    except Exception as e:
+        st.error(f"Error processing image: {e}")
+        return None
+
+def download_youtube_video(youtube_url):
+    try:
+        temp_dir = tempfile.gettempdir()
+        output_path = os.path.join(temp_dir, 'downloaded_video.mp4')
+        ydl_opts = {
+            'format': 'best',
+            'outtmpl': output_path
+        }
+        with youtube_dl.YoutubeDL(ydl_opts) as ydl:
+            ydl.download([youtube_url])
+        return output_path
+    except Exception as e:
+        st.error(f"Failed to retrieve video from YouTube. Error: {e}")
+        return None
+
+def process_video(input_video_path, output_video_path, model):
+    cap = cv2.VideoCapture(input_video_path)
+    if not cap.isOpened():
+        st.error(f"Error: Cannot open video file {input_video_path}")
+        return None
+
+    frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
+    frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
+    fps = cap.get(cv2.CAP_PROP_FPS)
+    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
+
+    out = cv2.VideoWriter(output_video_path, fourcc, fps, (frame_width, frame_height))
+
+    while True:
+        ret, frame = cap.read()
+        if not ret:
+            break
+
+        processed_frame = detect_objects(frame, model)
+        out.write(processed_frame)
+
+    cap.release()
+    out.release()
+
+    return output_video_path
+
+st.title("YOLOv8 Object Detection on Images and Videos")
+
+model_choice = st.selectbox("Select Model", ["yolov8s.pt", "yolov8m.pt"])
+
+tabs = st.tabs(["Image Detection", "Video Detection"])
+
+with tabs[0]:
+    st.header("Image Detection")
+    input_choice = st.radio("Select Input Method", ["Upload", "URL"])
+
+    if input_choice == "Upload":
+        uploaded_image = st.file_uploader("Upload Image", type=["jpg", "jpeg", "png"])
+        if uploaded_image is not None:
+            image = Image.open(uploaded_image)
+            processed_image = process_image_with_yolov8(model_choice, image=image)
+            if processed_image is not None:
+                st.image(processed_image, caption="Processed Image", use_column_width=True)
+    elif input_choice == "URL":
+        image_url = st.text_input("Image URL")
+        if image_url:
+            processed_image = process_image_with_yolov8(model_choice, url=image_url)
+            if processed_image is not None:
+                st.image(processed_image, caption="Processed Image", use_column_width=True)
+
+with tabs[1]:
+    st.header("Video Detection")
+    video_choice = st.radio("Select Input Method", ["Upload", "YouTube"])
+
+    if video_choice == "Upload":
+        uploaded_video = st.file_uploader("Upload Local Video", type=["mp4", "mov", "avi"])
+        if uploaded_video is not None:
+            input_video_path = os.path.join(tempfile.gettempdir(), uploaded_video.name)
+            with open(input_video_path, "wb") as f:
+                f.write(uploaded_video.read())
+            model = load_yolov8_model(model_choice)
+            output_video_path = os.path.join(tempfile.gettempdir(), "processed_video.mp4")
+            processed_video = process_video(input_video_path, output_video_path, model)
+            if processed_video is not None:
+                st.video(processed_video)
+    elif video_choice == "YouTube":
+        video_url = st.text_input("YouTube Video URL")
+        if video_url:
+            input_video_path = download_youtube_video(video_url)
+            if input_video_path:
+                model = load_yolov8_model(model_choice)
+                output_video_path = os.path.join(tempfile.gettempdir(), "processed_video.mp4")
+                processed_video = process_video(input_video_path, output_video_path, model)
+                if processed_video is not None:
+                    st.video(processed_video)