import os import cv2 import tempfile import requests import base64 import numpy as np import logging from dataclasses import dataclass from typing import Optional, Union, Tuple from PIL import Image from io import BytesIO from ultralytics import YOLO import streamlit as st import yt_dlp as youtube_dl from config import Config # Configure logging logging.basicConfig(level=logging.INFO) logger = logging.getLogger(__name__) @dataclass class DetectionResult: """Data class to store detection results""" success: bool image: Optional[np.ndarray] = None error_message: Optional[str] = None class YOLOModel: """Class to handle YOLO model operations""" def __init__(self, model_name: str = Config.DEFAULT_MODEL): self.model = self._load_model(model_name) def _load_model(self, model_name: str) -> Optional[YOLO]: """Load YOLO model with error handling""" try: return YOLO(model_name) except Exception as e: logger.error(f"Error loading model: {e}") return None def detect_objects(self, image: np.ndarray) -> DetectionResult: """Perform object detection on the input image""" if self.model is None: return DetectionResult(False, error_message="Model not loaded") try: results = self.model(image) annotated_image = image.copy() for result in results[0].boxes: x1, y1, x2, y2 = map(int, result.xyxy[0]) label = self.model.names[int(result.cls)] confidence = result.conf.item() if confidence < Config.CONFIDENCE_THRESHOLD: continue cv2.rectangle( annotated_image, (x1, y1), (x2, y2), Config.BBOX_COLOR, 2 ) label_text = f'{label} {confidence:.2f}' cv2.putText( annotated_image, label_text, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, Config.FONT_SCALE, Config.BBOX_COLOR, Config.FONT_THICKNESS ) return DetectionResult(True, annotated_image) except Exception as e: logger.error(f"Error during object detection: {e}") return DetectionResult(False, error_message=str(e)) class ImageProcessor: """Class to handle image processing operations""" def __init__(self, model: YOLOModel): self.model = model def process_image(self, image: Union[Image.Image, str]) -> DetectionResult: """Process image from various sources (PIL Image or URL)""" try: if isinstance(image, str): image = self._load_image_from_url(image) if image is None: return DetectionResult(False, error_message="Failed to load image") np_image = np.array(image) return self.model.detect_objects(np_image) except Exception as e: logger.error(f"Error processing image: {e}") return DetectionResult(False, error_message=str(e)) def _load_image_from_url(self, url: str) -> Optional[Image.Image]: """Load image from URL with support for base64""" try: if url.startswith('data:image'): header, encoded = url.split(',', 1) image_data = base64.b64decode(encoded) return Image.open(BytesIO(image_data)) else: response = requests.get(url) response.raise_for_status() return Image.open(BytesIO(response.content)) except Exception as e: logger.error(f"Error loading image from URL: {e}") return None class VideoProcessor: """Class to handle video processing operations""" def __init__(self, model: YOLOModel): self.model = model os.makedirs(Config.TEMP_DIR, exist_ok=True) def process_video(self, input_path: str) -> Tuple[bool, Optional[str]]: """Process video file and return path to processed video""" try: cap = cv2.VideoCapture(input_path) if not cap.isOpened(): return False, "Cannot open video file" output_path = os.path.join(Config.TEMP_DIR, "processed_video.mp4") self._setup_video_writer(cap, output_path) while True: ret, frame = cap.read() if not ret: break result = self.model.detect_objects(frame) if result.success: self.writer.write(result.image) cap.release() self.writer.release() return True, output_path except Exception as e: logger.error(f"Error processing video: {e}") return False, str(e) def _setup_video_writer(self, cap: cv2.VideoCapture, output_path: str): """Set up video writer with input video properties""" 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(*Config.VIDEO_OUTPUT_FORMAT) self.writer = cv2.VideoWriter( output_path, fourcc, fps, (frame_width, frame_height) ) def download_youtube_video(youtube_url: str) -> Optional[str]: """Download YouTube video and return path to downloaded file""" 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: logger.error(f"Failed to retrieve video from YouTube: {e}") return None def main(): """Main application function""" st.title("MULTIMEDIA OBJECT DETECTION USING YOLO") # Model selection with description st.subheader("Model Selection") model_choice = st.selectbox( "Select YOLO Model", options=Config.AVAILABLE_MODELS, index=Config.AVAILABLE_MODELS.index(Config.DEFAULT_MODEL), format_func=lambda x: f"{x} - {Config.YOLO_MODELS[x]}" ) # Display model capabilities model_type = "Detection" if "pose" in model_choice: model_type = "Pose Estimation" st.info("This model will detect and estimate human poses in the image/video.") elif "seg" in model_choice: model_type = "Instance Segmentation" st.info("This model will perform instance segmentation, creating precise masks for detected objects.") else: st.info("This model will detect and classify objects with bounding boxes.") # Initialize model and processors model = YOLOModel(model_choice) image_processor = ImageProcessor(model) video_processor = VideoProcessor(model) 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=Config.ALLOWED_IMAGE_TYPES ) if uploaded_image is not None: image = Image.open(uploaded_image) result = image_processor.process_image(image) if result.success: st.image(result.image, caption="Processed Image", use_container_width=True) else: st.error(result.error_message) elif input_choice == "URL": image_url = st.text_input("Image URL") if image_url: result = image_processor.process_image(image_url) if result.success: st.image(result.image, caption="Processed Image", use_container_width=True) else: st.error(result.error_message) 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=Config.ALLOWED_VIDEO_TYPES ) if uploaded_video is not None: input_video_path = os.path.join(Config.TEMP_DIR, uploaded_video.name) with open(input_video_path, "wb") as f: f.write(uploaded_video.read()) success, result = video_processor.process_video(input_video_path) if success: st.video(result) else: st.error(result) 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: success, result = video_processor.process_video(input_video_path) if success: st.video(result) else: st.error(result) if __name__ == "__main__": main()