Update backend/main.py

backend/main.py

@@ -2,21 +2,18 @@ from __future__ import annotations
 
 import os
 import tempfile
-from typing import List, Optional, Tuple
+import traceback
+from typing import List
 
 import numpy as np
 from fastapi import FastAPI, File, UploadFile, HTTPException
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.staticfiles import StaticFiles
-from fastapi.responses import FileResponse, JSONResponse
 
 import cv2  # type: ignore
 import joblib  # type: ignore
 import tensorflow as tf  # type: ignore
 from tensorflow.keras.applications.resnet import preprocess_input
-import os
-os.environ["SM_FRAMEWORK"] = "tf.keras"
-
 
 
 # -----------------------------
@@ -24,10 +21,10 @@ os.environ["SM_FRAMEWORK"] = "tf.keras"
 # -----------------------------
 MODEL_PATH = os.getenv("MODEL_PATH", "/app/model.keras")
 SCALER_PATH = os.getenv("SCALER_PATH", "/app/scaler.save")
+STATIC_DIR = os.getenv("STATIC_DIR", "/app/static")
 
 _model = None
 _scaler = None
-_preprocess_input = None
 
 
 def get_model():
@@ -39,7 +36,7 @@ def get_model():
                 "Place your Keras model at /app/model.keras or set MODEL_PATH."
             )
         # compile=False is safer for inference-only deployments
-        _model = tf.keras.models.load_model("/app/model.keras", compile=False)
+        _model = tf.keras.models.load_model(MODEL_PATH, compile=False)
     return _model
 
 
@@ -55,39 +52,34 @@ def get_scaler():
     return _scaler
 
 
-def get_preprocess_input():
-    global _preprocess_input
-    if _preprocess_input is None:
-        _preprocess_input = sm.get_preprocessing("resnet101")
-    return _preprocess_input
-
-
 # -----------------------------
 # Preprocessing (as requested)
 # -----------------------------
-def load_data(image_path):
+def load_data(image_path: str) -> tf.Tensor:
     image = tf.io.read_file(image_path)
     image = tf.io.decode_png(image, channels=3)
     image = tf.image.resize(image, [224, 224], method="bilinear")
     image = tf.cast(image, tf.float32)
-    image = preprocess_input(image)
+    image = preprocess_input(image)  # ResNet preprocessing
     return image
 
 
-
+# -----------------------------
+# Video -> frames
+# -----------------------------
 def extract_frames_to_pngs(video_bytes: bytes, max_frames: int = 300) -> List[str]:
     """Decode video bytes with OpenCV and write frames as PNGs to a temp dir.
-
     Returns a list of PNG file paths.
     """
     tmpdir = tempfile.mkdtemp(prefix="frames_")
     video_path = os.path.join(tmpdir, "input.mp4")
+
     with open(video_path, "wb") as f:
         f.write(video_bytes)
 
     cap = cv2.VideoCapture(video_path)
     if not cap.isOpened():
-        raise ValueError("Could not open uploaded video.")
+        raise ValueError("Could not open uploaded video. (Unsupported codec/container?)")
 
     paths: List[str] = []
     idx = 0
@@ -96,8 +88,6 @@ def extract_frames_to_pngs(video_bytes: bytes, max_frames: int = 300) -> List[st
         if not ok:
             break
 
-        # Ensure RGB->BGR handling: OpenCV reads BGR; we'll write PNG in BGR which is fine,
-        # because tf.io.decode_png reads the encoded pixels and we treat them as 3 channels.
         out_path = os.path.join(tmpdir, f"frame_{idx:05d}.png")
         cv2.imwrite(out_path, frame)
         paths.append(out_path)
@@ -110,14 +100,16 @@ def extract_frames_to_pngs(video_bytes: bytes, max_frames: int = 300) -> List[st
     return paths
 
 
+# -----------------------------
+# Post-processing helpers
+# -----------------------------
 def moving_average(x: np.ndarray, window: int = 7) -> np.ndarray:
     if window <= 1:
         return x
-    window = min(window, x.shape[0])
+    window = int(max(1, min(window, x.shape[0])))
     kernel = np.ones(window, dtype=np.float32) / float(window)
-    # pad to keep length
     pad = window // 2
-    xpad = np.pad(x, (pad, pad), mode="edge")
+    xpad = np.pad(x.astype(np.float32), (pad, pad), mode="edge")
     return np.convolve(xpad, kernel, mode="valid")
 
 
@@ -128,30 +120,69 @@ def compute_ef(edv: float, esv: float) -> float:
 
 
 def classify_heart_function(ef: float) -> str:
-    """Thresholds for EF-based function.
-
-    You asked for best thresholds; without calibration data, we use standard clinical cutoffs:
-    - Normal: >= 55
-    - Mild dysfunction: 40–54
-    - Heart failure: < 40
-    """
     if not np.isfinite(ef):
         return "heart failure"
-
     if ef >= 55.0:
         return "normal"
     if ef >= 40.0:
-        # Frontend expects this exact string union.
         return "mildly dysfunction"
     return "heart failure"
 
 
+def _normalize_model_output(raw, n_frames: int) -> np.ndarray:
+    """
+    Normalize model.predict output to shape (N, 1) float array suitable for scaler.inverse_transform.
+    Handles models that return:
+      - single array: (N,), (N,1), (N,k), (N,1,1), etc.
+      - list/tuple of arrays (multi-output)
+    """
+    # Multi-output model: choose the output whose first dim matches number of frames.
+    if isinstance(raw, (list, tuple)):
+        shapes = [np.asarray(x).shape for x in raw]
+        print("PRED LIST SHAPES:", shapes)
+
+        chosen = None
+        for r in raw:
+            r_arr = np.asarray(r)
+            if r_arr.ndim >= 1 and r_arr.shape[0] == n_frames:
+                chosen = r_arr
+                break
+        if chosen is None:
+            chosen = np.asarray(raw[0])
+        raw_arr = chosen
+    else:
+        raw_arr = np.asarray(raw)
+        print("PRED SHAPE:", raw_arr.shape)
+
+    raw_arr = np.asarray(raw_arr)
+
+    # Force to (N, 1)
+    if raw_arr.ndim == 1:
+        raw_arr = raw_arr.reshape(-1, 1)
+    elif raw_arr.ndim == 2:
+        if raw_arr.shape[0] != n_frames:
+            # Sometimes outputs come as (1, N) — fix that
+            if raw_arr.shape[1] == n_frames:
+                raw_arr = raw_arr.T
+        # If multiple columns, pick the first by default
+        if raw_arr.shape[1] != 1:
+            raw_arr = raw_arr[:, :1]
+    else:
+        # Flatten everything but the frame dimension
+        raw_arr = raw_arr.reshape(raw_arr.shape[0], -1)
+        raw_arr = raw_arr[:, :1]
+
+    if raw_arr.shape[0] != n_frames:
+        raise ValueError(f"Prediction length mismatch: got {raw_arr.shape[0]} but expected {n_frames}")
+
+    return raw_arr.astype(np.float32)
+
+
 # -----------------------------
 # FastAPI app
 # -----------------------------
 app = FastAPI()
 
-# Same-origin by default; allow all just in case spaces routes differ
 app.add_middleware(
     CORSMiddleware,
     allow_origins=["*"],
@@ -160,38 +191,34 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
-STATIC_DIR = os.getenv("STATIC_DIR", "/app/static")
-
 
 @app.post("/api/analyze")
 async def analyze(video: UploadFile = File(...)):
-    if video.content_type is None or "video" not in video.content_type:
-        # Some browsers may omit; still attempt
-        pass
-
     video_bytes = await video.read()
     if not video_bytes:
         raise HTTPException(status_code=400, detail="Empty video upload.")
 
     try:
-        frame_paths = extract_frames_to_pngs(video_bytes)
-    except Exception as e:
-        raise HTTPException(status_code=400, detail=f"Video decode error: {e}")
+        frame_paths = extract_frames_to_pngs(
+            video_bytes,
+            max_frames=int(os.getenv("MAX_FRAMES", "300")),
+        )
 
-    try:
         # Build a batch tensor [N, 224, 224, 3]
         batch = tf.stack([load_data(p) for p in frame_paths], axis=0)
+
         model = get_model()
-        preds = model.predict(batch, verbose=0)
+        raw_preds = model.predict(batch, verbose=0)
 
-        preds_np = np.array(preds)
-        preds_np = preds_np.reshape(-1, 1)
+        preds_np = _normalize_model_output(raw_preds, n_frames=batch.shape[0])
 
         scaler = get_scaler()
-        values = scaler.inverse_transform(preds_np).reshape(-1)
+        values = scaler.inverse_transform(preds_np).reshape(-1).astype(np.float32)
 
         # Smooth
-        smooth = moving_average(values, window=int(os.getenv("SMOOTH_WINDOW", "7")))
+        smooth_window = int(os.getenv("SMOOTH_WINDOW", "7"))
+        smooth = moving_average(values, window=smooth_window)
+
         edv = float(np.max(smooth))
         esv = float(np.min(smooth))
         ef = compute_ef(edv, esv)
@@ -202,11 +229,13 @@ async def analyze(video: UploadFile = File(...)):
             "heartFunction": heart_fn,
             "edv": round(edv, 2),
             "esv": round(esv, 2),
-            "numFrames": int(len(values)),
+            "numFrames": int(values.shape[0]),
         }
-    except RuntimeError as e:
-        raise HTTPException(status_code=500, detail=str(e))
+
+    except HTTPException:
+        raise
     except Exception as e:
+        print("ANALYZE ERROR TRACEBACK:\n", traceback.format_exc())
         raise HTTPException(status_code=500, detail=f"Inference error: {e}")