Update app.py

app.py

@@ -1,152 +1,52 @@
 import torch
-from transformers import AutoFeatureExtractor, AutoModelForImageClassification
+from torchvision import transforms
 from PIL import Image
-import io
-import base64
-import numpy as np
-from flask import Flask, request, jsonify
+import json
+import streamlit as st
 
-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}")
+# Charger les noms des classes
+with open("class_names.json", "r") as f:
+    class_names = json.load(f)
 
-# Flask API Setup
-app = Flask(__name__)
+# Charger le modèle
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+model = torch.load("efficientnet_b7_best.pth", map_location=device)
+model.eval()  # Mode évaluation
 
-# Initialize classifier (can be changed to any model)
-classifier = HuggingFaceClassifier(
-    model_name="microsoft/resnet-50"
-)
+# Définir la taille de l'image
+image_size = (224, 224)
 
-@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
+# Transformation pour l'image
+class GrayscaleToRGB:
+    def __call__(self, img):
+        return img.convert("RGB")
 
-@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)
-    })
+valid_test_transforms = transforms.Compose([
+    transforms.Grayscale(num_output_channels=1),
+    transforms.Resize(image_size),
+    GrayscaleToRGB(),  # Conversion en RGB
+    transforms.ToTensor(),
+    transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
+])
 
-@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)
-    })
+# Fonction de prédiction
+def predict_image(image):
+    image_tensor = valid_test_transforms(image).unsqueeze(0).to(device)
+    with torch.no_grad():
+        outputs = model(image_tensor)
+        _, predicted_class = torch.max(outputs, 1)
+    predicted_label = class_names[predicted_class.item()]
+    return predicted_label
 
-if __name__ == '__main__':
-    app.run(host='0.0.0.0', port=5000, debug=True)
+# Interface Streamlit
+st.title("Prédiction d'images avec PyTorch")
+st.write("Chargez une image pour obtenir une prédiction de classe.")
+
+uploaded_image = st.file_uploader("Téléchargez une image", type=["jpg", "jpeg", "png"])
+
+if uploaded_image is not None:
+    image = Image.open(uploaded_image)
+    st.image(image, caption="Image téléchargée", use_column_width=True)
+    
+    predicted_label = predict_image(image)
+    st.write(f"Prédiction de la classe : {predicted_label}")