Update app.py

app.py

@@ -10,127 +10,71 @@ import numpy as np
 import logging
 import sys
 import os
-from typing import Optional, Tuple
+from typing import Optional
 
-# ============================================================================
-# LOGGING CONFIGURATION
-# ============================================================================
+# 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)
-    ]
+    format='%(asctime)s - %(levelname)s - %(message)s',
+    handlers=[logging.StreamHandler(sys.stdout)]
 )
 logger = logging.getLogger(__name__)
 
-# ============================================================================
-# CONFIGURATION
-# ============================================================================
+# 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
+# Load the model
+try:
+    if os.path.exists(MODEL_PATH):
         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
-
+    else:
+        logger.error(f"❌ Model file not found at {MODEL_PATH}")
+        model = None
+except Exception as e:
+    logger.error(f"❌ Error loading model: {e}")
+    model = 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
-    """
+def classify_animal(image):
+    if image is None:
+        return "Please upload an image."
     
+    if model is None:
+        return "Model not loaded. Check server logs."
+
     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")
+        # Run inference
+        results = model(image)
         
-        # ========== MODEL AVAILABILITY CHECK ==========
-        if model is None:
-            logger.error("Model is not loaded")
-            return "❌ Critical Error: Model not loaded. Please restart the application."
+        # YOLOv8 classification returns a list of results
+        # We take the top prediction from the first result
+        result = results[0]
         
-        # ========== 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")
+        if result.probs is not None:
+            # Get index of the highest probability
+            top1_idx = result.probs.top1
+            conf = result.probs.top1conf.item()
+            label = result.names[top1_idx]
             
-            elif isinstance(image, Image.Image):
-                image_pil = image
-                lo
+            return f"Prediction: {label.upper()} ({conf:.2%})"
+        else:
+            return "No animals detected or classification failed."
+
+    except Exception as e:
+        logger.error(f"Inference error: {e}")
+        return f"Error during classification: {str(e)}"
+
+# Gradio Interface
+demo = gr.Interface(
+    fn=classify_animal,
+    inputs=gr.Image(type="pil", label="Upload Animal Image"),
+    outputs=gr.Textbox(label="Result"),
+    title="🐾 Animal Type Classifier",
+    description="Upload a photo of a butterfly, chicken, elephant, horse, spider, or squirrel to identify it.",
+    examples=[["example_elephant.jpg"]] if os.path.exists("example_elephant.jpg") else None,
+    cache_examples=False
+)
+
+if __name__ == "__main__":
+    demo.launch(server_name="0.0.0.0", server_port=7860)