Spaces:
Build error
Build error
File size: 4,882 Bytes
aac8470 8047dc3 395b902 8047dc3 aac8470 8047dc3 aac8470 8047dc3 aac8470 f996c5a aac8470 f996c5a aac8470 f996c5a 6614d3d aac8470 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 | """
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
|