Complete rewrite: Clean, simple backend for ConvNextV2-large-DogBreed

app.py

@@ -1,100 +1,73 @@
 """
-Production-ready Flask backend for deployment
-Optimized for Hugging Face Spaces / Railway / Render
+Clean Pet Backend for Hugging Face Spaces
+Using ConvNextV2-large-DogBreed model
 """
 
 from flask import Flask, request, jsonify
-import io
-from PIL import Image
 from flask_cors import CORS
-import logging
+from PIL import Image
+import io
 import os
 
-# Set cache directory to /tmp for Hugging Face Spaces
+# Set cache directories FIRST
 os.environ['HF_HOME'] = '/tmp/.cache'
 os.environ['TRANSFORMERS_CACHE'] = '/tmp/.cache'
 os.environ['TORCH_HOME'] = '/tmp/.cache'
 
 app = Flask(__name__)
-CORS(app, resources={r"/predict_pet": {"origins": "*"}})
-
-# Logging configuration
-logging.basicConfig(
-    level=logging.INFO,
-    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
-)
-logger = logging.getLogger(__name__)
+CORS(app)
 
-# Model global variables
+# Global model variables
 model = None
-image_processor = None
+processor = None
 
 def load_model():
     """Load model on first request"""
-    global model, image_processor
+    global model, processor
     
     if model is not None:
         return
     
-    try:
-        import torch
-        from transformers import AutoImageProcessor, AutoModelForImageClassification
-        
-        logger.info("Loading ConvNextV2-large-DogBreed model...")
-        
-        model_name = "Pavarissy/ConvNextV2-large-DogBreed"
-        
-        # Detect device
-        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-        logger.info(f"Using device: {device}")
-        
-        # Load model
-        image_processor = AutoImageProcessor.from_pretrained(model_name)
-        model = AutoModelForImageClassification.from_pretrained(model_name)
-        model = model.to(device)
-        model.eval()
-        
-        logger.info(f"✓ Model loaded successfully on {device}")
-        
-    except Exception as e:
-        logger.error(f"Failed to load model: {e}")
-        raise
+    print("🔄 Loading ConvNextV2-large-DogBreed model...")
+    
+    from transformers import AutoImageProcessor, AutoModelForImageClassification
+    import torch
+    
+    model_name = "Pavarissy/ConvNextV2-large-DogBreed"
+    
+    processor = AutoImageProcessor.from_pretrained(model_name)
+    model = AutoModelForImageClassification.from_pretrained(model_name)
+    
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    model = model.to(device)
+    model.eval()
+    
+    print(f"✅ Model loaded on {device}")
 
 @app.route('/', methods=['GET'])
-def health_check():
-    """Health check endpoint"""
+def health():
     return jsonify({
         'status': 'healthy',
-        'service': 'Dog Breed Prediction API',
         'model': 'ConvNextV2-large-DogBreed',
-        'accuracy': '91.39%',
-        'version': '1.0.0'
+        'accuracy': '91.39%'
     })
 
 @app.route('/predict_pet', methods=['POST'])
 def predict_pet():
-    """Predict dog breed from uploaded image"""
-    
     try:
-        # Load model if not loaded
+        # Load model if needed
         load_model()
         
-        # Validate request
+        # Get image
         if 'image' not in request.files:
-            return jsonify({'error': 'No image file provided'}), 400
-            
-        file = request.files['image']
+            return jsonify({'error': 'No image provided'}), 400
         
-        # Read and validate image
-        image_bytes = file.read()
-        pil_image = Image.open(io.BytesIO(image_bytes))
-        
-        if pil_image.mode != 'RGB':
-            pil_image = pil_image.convert('RGB')
+        file = request.files['image']
+        image = Image.open(io.BytesIO(file.read())).convert('RGB')
         
-        # Make prediction
+        # Process and predict
         import torch
-        inputs = image_processor(images=pil_image, return_tensors="pt")
+        inputs = processor(image, return_tensors="pt")
         
         device = next(model.parameters()).device
         inputs = {k: v.to(device) for k, v in inputs.items()}
@@ -102,35 +75,32 @@ def predict_pet():
         with torch.no_grad():
             outputs = model(**inputs)
             logits = outputs.logits
-            
-        probs = torch.nn.functional.softmax(logits, dim=-1)[0].cpu()
-        top_5_probs, top_5_indices = torch.topk(probs, 5)
         
-        # Format results
-        top_5_breeds = []
-        for prob, idx in zip(top_5_probs, top_5_indices):
-            top_5_breeds.append({
-                'breed': model.config.id2label[idx.item()],
-                'confidence': float(prob.item())
-            })
+        # Get top prediction
+        probs = torch.nn.functional.softmax(logits, dim=-1)[0]
+        top_prob, top_idx = torch.max(probs, dim=0)
+        
+        predicted_breed = model.config.id2label[top_idx.item()]
+        confidence = float(top_prob.item())
         
-        logger.info(f"Prediction: {top_5_breeds[0]['breed']} ({top_5_breeds[0]['confidence']:.2%})")
+        print(f"✅ Predicted: {predicted_breed} ({confidence:.2%})")
         
+        # Return in format expected by frontend
         return jsonify({
-            'predicted_label': top_5_breeds[0]['breed'],
-            'breed': top_5_breeds[0]['breed'],
-            'confidence': top_5_breeds[0]['confidence'],
-            'top_5': top_5_breeds,
-            'model': 'ConvNextV2-large-DogBreed',
-            'accuracy': '91.39%',
+            'predicted_label': predicted_breed,
+            'breed': predicted_breed,
+            'confidence': confidence,
             'detection': {'box': {'x': 0, 'y': 0, 'width': 0, 'height': 0}},
             'gender': 'Unknown'
         })
         
     except Exception as e:
-        logger.error(f"Error: {str(e)}", exc_info=True)
+        print(f"❌ Error: {str(e)}")
+        import traceback
+        traceback.print_exc()
         return jsonify({'error': str(e)}), 500
 
 if __name__ == '__main__':
     port = int(os.environ.get('PORT', 7860))
+    print(f"🚀 Starting server on port {port}")
     app.run(host='0.0.0.0', port=port, debug=False)

requirements.txt

@@ -2,6 +2,6 @@ flask==3.0.0
 flask-cors==4.0.0
 transformers==4.35.0
 torch==2.1.0
+torchvision==0.16.0
 pillow==10.1.0
 accelerate==0.24.0
-