Update app.py

app.py

@@ -6,7 +6,7 @@ from huggingface_hub import snapshot_download
 import os
 import sys
 
-print("🦁 Setting up zoolang model...")
+print("🦁 Setting up zoolang model for dog emotion classification...")
 
 # Download your model files
 model_path = snapshot_download(repo_id="vsimmer/zoolang")
@@ -19,13 +19,17 @@ sys.path.insert(0, model_path)
 from model import get_model
 from handler import EndpointHandler
 
+# Your class names
+CLASS_NAMES = ['aggressive', 'anxious', 'frightened', 'happy', 'inquisitive']
+
 # Initialize your handler
 print("Initializing model handler...")
 handler = EndpointHandler(path=model_path)
 print("✅ Model loaded successfully!")
+print(f"✅ Classes: {CLASS_NAMES}")
 
 def classify_image(image):
-    """Classify an image using your zoolang model"""
+    """Classify dog emotion using your zoolang model"""
     if image is None:
         return {"No image": 1.0}
     
@@ -33,43 +37,4 @@ def classify_image(image):
         # Convert PIL image to bytes for your handler
         import io
         img_byte_arr = io.BytesIO()
-        image.save(img_byte_arr, format='PNG')
-        img_byte_arr.seek(0)
-        
-        # Use your handler
-        result = handler({"inputs": img_byte_arr})
-        
-        # Your handler returns {"label": prediction} where prediction is a class index
-        # Convert to probabilities (since we don't have softmax from your model)
-        # For now, return the predicted class with high confidence
-        if isinstance(result, dict) and "label" in result:
-            class_index = result["label"]
-            # Create mock probabilities - your model's class index as 100% confidence
-            # Since your model has 5 classes (0-4)
-            class_names = [f"Class_{i}" for i in range(5)]  # Replace with your actual class names
-            
-            # Return the predicted class as the top result
-            return {
-                class_names[class_index]: 0.95,  # High confidence for predicted class
-                "Other": 0.05
-            }
-        else:
-            return {"Unknown result format": 0.5}
-            
-    except Exception as e:
-        print(f"Classification error: {e}")
-        return {"Classification failed": 1.0}
-
-# Create Gradio interface
-iface = gr.Interface(
-    fn=classify_image,
-    inputs=gr.Image(type="pil", label="Upload Image"),
-    outputs=gr.Label(num_top_classes=5, label="Zoolang Prediction"),
-    title="🦁 Zoolang Image Classifier",
-    description="Using your custom MobileNet V2 model with 5 classes",
-    allow_flagging="never",
-    examples=None
-)
-
-if __name__ == "__main__":
-    iface.launch(server_name="0.0.0.0", server_port=7860)
+        image.save(img_byte_arr, format='PNG