Create app.py

app.py

@@ -0,0 +1,76 @@
+import os
+from typing import Any
+from flask import Flask, jsonify, request, send_from_directory
+from PIL import Image
+import torch
+import torch.nn.functional as F
+from dotenv import load_dotenv
+from model_loader import load_alexnet_model, preprocess_image
+
+load_dotenv(override=True)
+
+# HF sets PORT dynamically. Fall back to 7860 locally.
+PORT = int(os.getenv("PORT", os.getenv("FLASK_PORT", "7860")))
+HOST = "0.0.0.0"
+MODEL_PATH = os.getenv("MODEL_PATH", "models/alexnext_vsf_bext.pth")
+
+# Single worker is safest for GPU inference
+torch.set_num_threads(1)
+
+# Create app and static hosting
+app = Flask(__name__, static_folder="static", static_url_path="")
+
+# Device selection
+DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+# Load model once at startup
+model, classes = load_alexnet_model(MODEL_PATH, device=DEVICE)
+model.to(DEVICE).eval()
+
+@app.get("/")
+def root() -> Any:
+    # serve your frontend
+    return send_from_directory(app.static_folder, "index.html")
+
+@app.get("/health")
+def health() -> Any:
+    return jsonify({"status": "ok", "device": str(DEVICE)})
+
+def load_image(file_stream):
+    return Image.open(file_stream).convert("RGB")
+
+@app.post("/predict_AlexNet")
+def predict_alexnet() -> Any:
+    if "image" not in request.files:
+        return jsonify({"error": "Missing file field 'image'."}), 400
+
+    file = request.files["image"]
+    if not file:
+        return jsonify({"error": "Empty file."}), 400
+
+    try:
+        img = load_image(file.stream)
+        input_tensor = preprocess_image(img).to(DEVICE)
+
+        with torch.no_grad():
+            output = model(input_tensor)
+            probabilities = F.softmax(output[0], dim=0).detach().cpu()
+
+        pred_prob, pred_idx = torch.max(probabilities, dim=0)
+        predicted_class = classes[int(pred_idx)]
+
+        result = {
+            "class": predicted_class,
+            "confidence": float(pred_prob),
+            "probabilities": {
+                cls: float(prob) for cls, prob in zip(classes, probabilities.tolist())
+            },
+        }
+        return jsonify(result)
+
+    except Exception as e:
+        return jsonify({"error": f"Failed to process image: {e}"}), 400
+
+if __name__ == "__main__":
+    debug = bool(int(os.getenv("FLASK_DEBUG", "0")))
+    app.run(host=HOST, port=PORT, debug=debug)