Update app.py

app.py

@@ -0,0 +1,152 @@
+import torch
+from transformers import AutoFeatureExtractor, AutoModelForImageClassification
+from PIL import Image
+import io
+import base64
+import numpy as np
+from flask import Flask, request, jsonify
+
+class HuggingFaceClassifier:
+    def __init__(self, model_name="microsoft/resnet-50"):
+        """
+        Initialize Hugging Face model and feature extractor
+        
+        Args:
+            model_name (str): Hugging Face model identifier
+        """
+        try:
+            # Load pre-trained model and feature extractor
+            self.feature_extractor = AutoFeatureExtractor.from_pretrained(model_name)
+            self.model = AutoModelForImageClassification.from_pretrained(model_name)
+            
+            # Move to GPU if available
+            self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+            self.model.to(self.device)
+            self.model.eval()
+        
+        except Exception as e:
+            raise ValueError(f"Model loading error: {e}")
+    
+    def preprocess_image(self, image):
+        """
+        Preprocess image for model input
+        
+        Args:
+            image (PIL.Image): Input image
+        
+        Returns:
+            torch.Tensor: Preprocessed image tensor
+        """
+        # Preprocess image using feature extractor
+        inputs = self.feature_extractor(images=image, return_tensors="pt")
+        return inputs.pixel_values.to(self.device)
+    
+    def predict(self, image):
+        """
+        Predict image classification
+        
+        Args:
+            image (PIL.Image): Input image
+        
+        Returns:
+            list: Top prediction results
+        """
+        try:
+            # Preprocess image
+            inputs = self.preprocess_image(image)
+            
+            # Perform prediction
+            with torch.no_grad():
+                outputs = self.model(inputs)
+                logits = outputs.logits
+                probabilities = torch.softmax(logits, dim=-1)
+                top_k = torch.topk(probabilities, k=5)
+            
+            # Process results
+            predicted_classes = [
+                {
+                    "label": self.model.config.id2label[idx.item()],
+                    "score": prob.item()
+                }
+                for idx, prob in zip(top_k.indices[0], top_k.values[0])
+            ]
+            
+            return predicted_classes
+        
+        except Exception as e:
+            raise RuntimeError(f"Prediction error: {e}")
+
+# Flask API Setup
+app = Flask(__name__)
+
+# Initialize classifier (can be changed to any model)
+classifier = HuggingFaceClassifier(
+    model_name="microsoft/resnet-50"
+)
+
+@app.route('/predict', methods=['POST'])
+def predict_image():
+    """
+    Image classification endpoint
+    Supports base64 and file upload
+    """
+    try:
+        # Handle base64 encoded image
+        if 'image' in request.json:
+            image_data = base64.b64decode(request.json['image'])
+            image = Image.open(io.BytesIO(image_data))
+        
+        # Handle file upload
+        elif 'file' in request.files:
+            image = Image.open(request.files['file'])
+        
+        else:
+            return jsonify({
+                'error': 'No image provided',
+                'status': 'failed'
+            }), 400
+        
+        # Perform prediction
+        predictions = classifier.predict(image)
+        
+        return jsonify({
+            'predictions': predictions,
+            'status': 'success'
+        })
+    
+    except Exception as e:
+        return jsonify({
+            'error': str(e),
+            'status': 'failed'
+        }), 500
+
+@app.route('/models', methods=['GET'])
+def available_models():
+    """
+    List available pre-trained models
+    """
+    models = [
+        "microsoft/resnet-50",
+        "google/vit-base-patch16-224",
+        "facebook/vit-mae-base",
+        "microsoft/beit-base-patch16-224"
+    ]
+    
+    return jsonify({
+        'models': models,
+        'total_models': len(models)
+    })
+
+@app.route('/health', methods=['GET'])
+def health_check():
+    """
+    API health check endpoint
+    """
+    return jsonify({
+        'status': 'healthy',
+        'model': classifier.model.config.model_type,
+        'device': str(classifier.device)
+    })
+
+if __name__ == '__main__':
+    app.run(host='0.0.0.0', port=5000, debug=True)