Reuse local main.py + numpy 1.26.4

main.py

@@ -1,3 +1,7 @@
+from fastapi import FastAPI, File, UploadFile, HTTPException
+from fastapi.middleware.cors import CORSMiddleware
+from fastapi.responses import JSONResponse
+
 import os
 import logging
 from typing import List
@@ -6,111 +10,117 @@ import numpy as np
 from PIL import Image
 
 import tensorflow as tf
-from fastapi import FastAPI, File, UploadFile, HTTPException
-from fastapi.middleware.cors import CORSMiddleware
-from fastapi.responses import JSONResponse
-
 from huggingface_hub import hf_hub_download
 
 
-# -----------------------
-# Config
-# -----------------------
 logging.basicConfig(level=logging.INFO)
-logger = logging.getLogger("hf-space-fastapi")
-
-HF_REPO_ID = os.environ.get("HF_REPO_ID", "warresnaet/masterclass-2025")
-HF_MODEL_FILENAME = os.environ.get("HF_MODEL_FILENAME", "model.keras")
-HF_REVISION = os.environ.get("HF_REVISION", "main")
-LOCAL_MODEL_DIR = os.environ.get("LOCAL_MODEL_DIR", "./model")
-
-ANIMALS: List[str] = ["Cat", "Dog", "Panda"]
+logger = logging.getLogger(__name__)
 
 
-# -----------------------
-# App
-# -----------------------
-app = FastAPI(title="Animal Classification API (HF Space)")
+app = FastAPI(title="Animal Classification API")
 
 app.add_middleware(
     CORSMiddleware,
     allow_origins=["*"],
+    allow_credentials=True,
     allow_methods=["*"],
     allow_headers=["*"],
-    allow_credentials=True,
 )
 
 
-def _ensure_model_local() -> str:
-    """Download model file from HF if missing; return the absolute path to the model file."""
-    os.makedirs(LOCAL_MODEL_DIR, exist_ok=True)
-    local_model_path = os.path.join(LOCAL_MODEL_DIR, HF_MODEL_FILENAME)
-    if os.path.exists(local_model_path):
-        return os.path.abspath(local_model_path)
-
-    logger.info(
-        f"Downloading model from HF Hub: repo_id={HF_REPO_ID}, "
-        f"filename={HF_MODEL_FILENAME}, revision={HF_REVISION}"
-    )
-    downloaded = hf_hub_download(
-        repo_id=HF_REPO_ID,
-        filename=HF_MODEL_FILENAME,
-        repo_type="model",
-        revision=HF_REVISION,
-        local_dir=LOCAL_MODEL_DIR,
-    )
-    return os.path.abspath(downloaded)
-
-
-def _load_model() -> tf.keras.Model:
-    """Load the Keras model from the local path."""
-    model_path = _ensure_model_local()
-    logger.info(f"Loading Keras model from: {model_path}")
-    model = tf.keras.models.load_model(model_path)
-    logger.info("Model loaded")
-    return model
+ANIMALS: List[str] = ["Cat", "Dog", "Panda"]
 
 
-MODEL: tf.keras.Model = _load_model()
+# Resolve model path. Prefer an environment variable for flexibility. As a fallback
+# try a model.keras file in the repository root (one level up from this file).
+base_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
+model_path = os.environ.get("MODEL_PATH") or os.path.join(base_dir, "model.keras")
+
+
+model = None
+try:
+    if os.path.exists(model_path):
+        logger.info(f"Loading model from: {model_path}")
+        model = tf.keras.models.load_model(model_path)
+        logger.info("Model loaded successfully")
+    else:
+        # Try Hugging Face Hub fallback when local model is missing
+        repo_id = os.environ.get("HF_REPO_ID")
+        filename = os.environ.get("HF_MODEL_FILENAME", "model.keras")
+        revision = os.environ.get("HF_REVISION")
+        if repo_id:
+            try:
+                logger.info(
+                    f"Local model not found. Trying HF Hub: repo_id={repo_id}, filename={filename}, revision={revision}"
+                )
+                local_dir = os.path.join(base_dir, "hf_model")
+                os.makedirs(local_dir, exist_ok=True)
+                downloaded_path = hf_hub_download(
+                    repo_id=repo_id,
+                    filename=filename,
+                    repo_type="model",
+                    revision=revision,
+                    local_dir=local_dir,
+                )
+                logger.info(f"Downloaded model file to: {downloaded_path}")
+                model = tf.keras.models.load_model(downloaded_path)
+                logger.info("Model loaded successfully from HF Hub")
+            except Exception:
+                logger.exception("HF Hub fallback failed")
+        if model is None:
+            logger.warning(
+                "Model not available. Set MODEL_PATH or HF_REPO_ID (+ HF_MODEL_FILENAME) environment variables."
+            )
+except Exception as e:
+    logger.exception("Failed to load model")
+    model = None
 
 
 @app.get("/")
-def root():
-    return {"ok": True, "classes": ANIMALS}
-
-
-@app.get("/health")
-def health():
-    return {"status": "healthy", "model_loaded": MODEL is not None}
+def read_root():
+    return {"hello": "world"}
 
 
 @app.post("/upload/image")
 async def upload_image(img: UploadFile = File(...)):
+    """Accept an uploaded image, resize to (64,64), run model.predict and return the label.
+
+    If the model is not available the endpoint will return 503.
     """
-    Accept an image, resize to (64, 64), run prediction, and return label + scores.
-    """
-    if MODEL is None:
-        raise HTTPException(status_code=503, detail="Model is not loaded")
+    if model is None:
+        raise HTTPException(status_code=503, detail="Model is not loaded on the server")
 
     try:
-        image = Image.open(img.file).convert("RGB").resize((64, 64))
-        arr = np.array(image, dtype=np.float32)
-        batch = np.expand_dims(arr, axis=0)
-
-        probs = MODEL.predict(batch, verbose=0)[0]
-        probs = np.asarray(probs, dtype=np.float32)
-        idx = int(np.argmax(probs))
-        label = ANIMALS[idx] if 0 <= idx < len(ANIMALS) else str(idx)
+        # Read image bytes and ensure RGB
+        original_image = Image.open(img.file).convert("RGB")
+        # Preprocess the image
+        original_image = original_image.resize((64, 64))
+        # Training used raw pixel values [0-255], NOT normalized to [0-1]
+        img_array = np.array(original_image, dtype=np.float32)
+        img_array = np.expand_dims(img_array, axis=0)
+
+        predictions = model.predict(img_array)
+        # predictions might be shape (1, N)
+        probs = np.asarray(predictions).squeeze()
+        if probs.ndim == 0:
+            # Model returned a single value
+            label_idx = int(np.round(probs))
+        else:
+            label_idx = int(np.argmax(probs))
+
+        label = ANIMALS[label_idx] if 0 <= label_idx < len(ANIMALS) else str(label_idx)
 
         return JSONResponse({"label": label, "scores": probs.tolist()})
-    except Exception as e:
+
+    except Exception:
         logger.exception("Failed to process image")
-        raise HTTPException(status_code=400, detail=f"Bad image: {e}")
+        raise HTTPException(status_code=400, detail="Failed to process image")
 
 
 if __name__ == "__main__":
-    # For local testing. In HF Spaces (Docker SDK), the container will run uvicorn with this app.
+    # Run with: python main.py
+    # Use Uvicorn as the ASGI server. MODEL_PATH and PORT can be overridden via env vars.
     import uvicorn
 
-    port = int(os.environ.get("PORT", "7860"))
+    port = int(os.environ.get("PORT", 8000))
     uvicorn.run("main:app", host="0.0.0.0", port=port, reload=False)

requirements.txt

@@ -1,7 +1,7 @@
 fastapi==0.116.1
 uvicorn[standard]==0.23.2
 Pillow==10.1.0
-numpy<2.0.0
+numpy==1.26.4
 tensorflow-cpu==2.16.1
 huggingface_hub>=0.20.0
 python-multipart