app.py

"""
Animal Type Classification App
A robust Gradio application for classifying animals using YOLOv8
"""

import gradio as gr
from ultralytics import YOLO
from PIL import Image
import numpy as np
import logging
import sys
import os
from typing import Optional, Tuple

# ============================================================================
# LOGGING CONFIGURATION
# ============================================================================
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
    handlers=[
        logging.FileHandler('animal_classifier.log'),
        logging.StreamHandler(sys.stdout)
    ]
)
logger = logging.getLogger(__name__)

# ============================================================================
# CONFIGURATION
# ============================================================================
MODEL_PATH = "best_animal_classifier.pt"
CLASS_NAMES = ["butterfly", "chicken", "elephant", "horse", "spider", "squirrel"]
CONFIDENCE_THRESHOLD = 0.5
MIN_DETECTIONS = 1

# ============================================================================
# GLOBAL MODEL VARIABLE
# ============================================================================
model = None

# ============================================================================
# MODEL INITIALIZATION WITH EXCEPTION HANDLING
# ============================================================================
def load_model() -> Optional[YOLO]:
    """
    Load the YOLO model with comprehensive error handling.
    
    Returns:
        YOLO: Loaded model object or None if loading fails
    """
    global model
    
    try:
        logger.info(f"Attempting to load model from: {MODEL_PATH}")
        
        # Check if model file exists
        if not os.path.exists(MODEL_PATH):
            logger.error(f"Model file not found at: {MODEL_PATH}")
            raise FileNotFoundError(
                f"Model file '{MODEL_PATH}' does not exist. "
                f"Please ensure the file is in the correct location."
            )
        
        # Check if file has read permissions
        if not os.access(MODEL_PATH, os.R_OK):
            logger.error(f"No read permission for model file: {MODEL_PATH}")
            raise PermissionError(
                f"No read permission for model file: {MODEL_PATH}"
            )
        
        # Load the model
        model = YOLO(MODEL_PATH)
        logger.info("✅ Model loaded successfully!")
        
        return model
    
    except FileNotFoundError as e:
        logger.error(f"FileNotFoundError: {e}")
        return None
    
    except PermissionError as e:
        logger.error(f"PermissionError: {e}")
        return None
    
    except Exception as e:
        logger.error(f"Unexpected error loading model: {type(e).__name__}: {e}")
        return None


# ============================================================================
# CLASSIFICATION FUNCTION WITH ROBUST ERROR HANDLING
# ============================================================================
def classify_animal(image: Optional[np.ndarray]) -> str:
    """
    Classify an animal in the provided image using YOLOv8.
    
    Args:
        image (Optional[np.ndarray]): Input image as numpy array or PIL Image
    
    Returns:
        str: Classification result with confidence score or error message
    """
    
    try:
        # ========== INPUT VALIDATION ==========
        if image is None:
            logger.warning("No image provided")
            return "❌ Error: No image provided. Please upload an image."
        
        logger.info("Image received for classification")
        
        # ========== MODEL AVAILABILITY CHECK ==========
        if model is None:
            logger.error("Model is not loaded")
            return "❌ Critical Error: Model not loaded. Please restart the application."
        
        # ========== IMAGE TYPE CONVERSION ==========
        try:
            if isinstance(image, np.ndarray):
                # Validate numpy array dimensions
                if image.ndim not in [2, 3, 4]:
                    logger.error(f"Invalid image dimensions: {image.ndim}")
                    return "❌ Error: Invalid image dimensions. Expected 2D, 3D, or 4D array."
                
                # Validate data type
                if not np.issubdtype(image.dtype, np.integer):
                    logger.warning(f"Unexpected image dtype: {image.dtype}, attempting conversion")
                    image = image.astype('uint8')
                
                # Convert to PIL Image
                image_pil = Image.fromarray(image.astype('uint8'))
                logger.debug("Converted numpy array to PIL Image")
            
            elif isinstance(image, Image.Image):
                image_pil = image
                lo