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 @st.cache_resource def load_yolo_model(model_name: str) -> YOLO: """Load YOLO model with caching""" try: if model_name not in Config.AVAILABLE_MODELS: raise ValueError(f"Invalid model name: {model_name}") return YOLO(model_name) except Exception as e: logger.error(f"Error loading model: {e}") raise RuntimeError(f"Failed to load model: {str(e)}") class YOLOModel: """Class to handle YOLO model operations""" def __init__(self, model_name: str = Config.DEFAULT_MODEL): if not Config.validate_config(): raise RuntimeError("Invalid configuration") self.model = load_yolo_model(model_name) 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") # Convert image to RGB if it has an alpha channel if image.mode == 'RGBA': image = image.convert('RGB') 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""" cap = None writer = None progress_bar = st.progress(0) status_text = st.empty() try: cap = cv2.VideoCapture(input_path) if not cap.isOpened(): return False, "Cannot open video file" total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) if total_frames <= 0: return False, "Invalid video file" output_path = os.path.join(Config.TEMP_DIR, "processed_video.mp4") writer = self._setup_video_writer(cap, output_path) frame_count = 0 while True: ret, frame = cap.read() if not ret: break frame_count += 1 progress = min(100, int(frame_count * 100 / total_frames)) progress_bar.progress(progress) status_text.text(f"Processing frame {frame_count}/{total_frames}") result = self.model.detect_objects(frame) if result.success: writer.write(result.image) return True, output_path except Exception as e: logger.error(f"Error processing video: {e}") return False, str(e) finally: if cap is not None: cap.release() if writer is not None: writer.release() progress_bar.empty() status_text.empty() def _setup_video_writer(self, cap: cv2.VideoCapture, output_path: str) -> cv2.VideoWriter: """Set up video writer with input video properties""" width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) fps = cap.get(cv2.CAP_PROP_FPS) fourcc = cv2.VideoWriter_fourcc(*Config.VIDEO_OUTPUT_FORMAT) return cv2.VideoWriter(output_path, fourcc, fps, (width, 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 cleanup_temp_files(): """Clean up temporary files""" try: for file in os.listdir(Config.TEMP_DIR): file_path = os.path.join(Config.TEMP_DIR, file) try: if os.path.isfile(file_path): os.unlink(file_path) except Exception as e: logger.error(f"Error deleting {file_path}: {e}") except Exception as e: logger.error(f"Error cleaning up temp directory: {e}") def validate_image(image: Image.Image) -> Tuple[bool, str]: """Validate image format and properties""" try: # Check image mode if image.mode not in ['RGB', 'RGBA']: return False, f"Unsupported image mode: {image.mode}" # Check image size max_dimension = 1920 width, height = image.size if width > max_dimension or height > max_dimension: return False, f"Image too large. Maximum dimension: {max_dimension}px" # Check if image is valid image.verify() return True, "Image is valid" except Exception as e: return False, str(e) 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]}" ) # Initialize model using session state if 'model' not in st.session_state or st.session_state.get('model_choice') != model_choice: try: st.session_state.model = YOLOModel(model_choice) st.session_state.model_choice = model_choice except Exception as e: st.error(f"Error loading model: {str(e)}") return model = st.session_state.model image_processor = ImageProcessor(model) video_processor = VideoProcessor(model) # 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.") 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": try: uploaded_video = st.file_uploader( "Upload Local Video", type=Config.ALLOWED_VIDEO_TYPES ) if uploaded_video is not None: if uploaded_video.size > 200 * 1024 * 1024: # 200MB limit st.error("Video file is too large. Please upload a file smaller than 200MB.") return 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()) try: success, result = video_processor.process_video(input_video_path) if success: st.video(result) else: st.error(f"Error processing video: {result}") finally: cleanup_temp_files() except Exception as e: st.error(f"An error occurred: {str(e)}") elif video_choice == "YouTube": video_url = st.text_input("YouTube Video URL") if video_url: with st.spinner("Downloading video..."): input_video_path = download_youtube_video(video_url) if input_video_path: try: success, result = video_processor.process_video(input_video_path) if success: st.video(result) else: st.error(result) finally: cleanup_temp_files() else: st.error("Failed to download YouTube video") if __name__ == "__main__": main()