Update model.py

model.py

@@ -1,66 +1,40 @@
-from flask import Flask, request, jsonify
-from flask_cors import CORS
-from PIL import Image
 import torch
-from torchvision import transforms, models
-
-# Initialize Flask app
-app = Flask(__name__)
-
-# Enable CORS
-CORS(app)
-
-# Load the trained model
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-
-# Define the model architecture
-model = models.resnet152()
-model.fc = torch.nn.Linear(model.fc.in_features, 26)  # Adjust for the number of classes
-model.load_state_dict(torch.load("trained_model.pth", map_location=device))
-model = model.to(device)
-model.eval()
-
-# Define preprocessing for the input image
-preprocess = transforms.Compose([
-    transforms.Resize((224, 224)),
-    transforms.ToTensor(),
-    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
-])
-
-# Class labels (replace with your dataset's classes)
-CLASS_LABELS = [
-"bluebell", "buttercup", "colts_foot", "corn_poppy", "cowslip", 
-"crocus", "daffodil", "daisy", "dandelion", "foxglove", 
-"fritillary", "geranium", "hibiscus", "iris", "lily_valley", 
-"pansy", "petunia", "rose", "snowdrop", "sunflower", 
-"tigerlily", "tulip", "wallflower", "water_lily", "wild_tulip", 
-"windflower"
-]
-
-@app.route("/predict", methods=["POST"])
-def predict():
-    if "file" not in request.files:
-        return jsonify({"error": "No file uploaded"}), 400
-
-    file = request.files["file"]
-
-    try:
-        # Load and preprocess the image
-        image = Image.open(file.stream).convert("RGB")
-        input_tensor = preprocess(image).unsqueeze(0).to(device)
-
-        # Perform inference
+from torchvision import models, transforms
+from PIL import Image
+import json
+
+# Load model
+class CustomResNet:
+    def __init__(self, model_path, num_classes):
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        self.model = models.resnet152(pretrained=False)
+        self.model.fc = torch.nn.Linear(self.model.fc.in_features, num_classes)
+        self.model.load_state_dict(torch.load(model_path, map_location=self.device))
+        self.model.to(self.device)
+        self.model.eval()
+
+        # Preprocessing
+        self.preprocess = transforms.Compose([
+            transforms.Resize((224, 224)),
+            transforms.ToTensor(),
+            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
+        ])
+
+    def predict(self, image_bytes):
+        # Load and preprocess image
+        image = Image.open(image_bytes).convert("RGB")
+        tensor = self.preprocess(image).unsqueeze(0).to(self.device)
+
+        # Make prediction
         with torch.no_grad():
-            outputs = model(input_tensor)
-            _, predicted_class = torch.max(outputs, 1)
-        
-        predicted_label = CLASS_LABELS[predicted_class.item()]
-        print(f"Predicted class: {predicted_label}")
-        return jsonify({"predicted_class": predicted_label})
+            outputs = self.model(tensor)
+            _, predicted = torch.max(outputs, 1)
+
+        return predicted.item()
 
-    except Exception as e:
-        return jsonify({"error": f"Error during prediction: {str(e)}"}), 500
 
-# Run the app (Hugging Face Spaces requires `app.run()` here)
-if __name__ == "__main__":
-    app.run(host="0.0.0.0", port=8080)
+# API function
+def load_model():
+    with open("config.json", "r") as f:
+        config = json.load(f)
+    return CustomResNet("trained_model.pth", config["num_labels"])