Add initial project files including Dockerfile, main application logic, model checkpoint, and requirements

.gitignore

@@ -0,0 +1,2 @@
+__pycache__/
+*.pyc

Dockerfile

@@ -0,0 +1,15 @@
+FROM python:3.10-slim
+
+RUN apt-get update \
+    && apt-get install -y --no-install-recommends ffmpeg \
+    && rm -rf /var/lib/apt/lists/*
+
+WORKDIR /app
+
+COPY requirements.txt .
+RUN pip install --no-cache-dir -r requirements.txt
+
+COPY . .
+
+ENV PORT=7860
+CMD ["bash", "-lc", "uvicorn main:app --host 0.0.0.0 --port ${PORT}"]

elephant_stress_r2plus1d18.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a4db9df704d55a96b31c538ed6e2c54bb1fff38eebfaaac6596998abc55c6c7a
+size 125385517

main.py

@@ -0,0 +1,143 @@
+import os
+from tempfile import NamedTemporaryFile
+from typing import Dict, List
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torchvision.transforms as T
+from decord import VideoReader, cpu
+from fastapi import FastAPI, File, HTTPException, UploadFile
+from fastapi.middleware.cors import CORSMiddleware
+from torchvision.models.video import R2Plus1D_18_Weights, r2plus1d_18
+
+app = FastAPI(title="Elephant Stress API")
+
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_credentials=False,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+CKPT_PATH = "elephant_stress_r2plus1d18.pt"
+CLIP_LEN = 16
+NUM_CLASSES = 2
+
+EVAL_FRAME_TRANSFORM = T.Compose(
+    [
+        T.ToPILImage(),
+        T.Resize((128, 171)),
+        T.CenterCrop((112, 112)),
+        T.ToTensor(),
+    ]
+)
+
+KIN_MEAN = torch.tensor([0.43216, 0.394666, 0.37645]).view(3, 1, 1)
+KIN_STD = torch.tensor([0.22803, 0.22145, 0.216989]).view(3, 1, 1)
+
+
+def _normalize_state_dict(loaded_obj: object) -> Dict[str, torch.Tensor]:
+    if isinstance(loaded_obj, dict):
+        for key in ("state_dict", "model_state_dict", "model"):
+            nested = loaded_obj.get(key)
+            if isinstance(nested, dict):
+                loaded_obj = nested
+                break
+    if not isinstance(loaded_obj, dict):
+        raise ValueError("Checkpoint does not contain a valid state dict.")
+
+    normalized: Dict[str, torch.Tensor] = {}
+    for key, value in loaded_obj.items():
+        clean_key = key[7:] if key.startswith("module.") else key
+        normalized[clean_key] = value
+    return normalized
+
+
+def _load_model() -> torch.nn.Module:
+    if not os.path.exists(CKPT_PATH):
+        raise FileNotFoundError(
+            f"Missing checkpoint: {CKPT_PATH}. Add it to repo root with git-lfs."
+        )
+
+    model = r2plus1d_18(weights=R2Plus1D_18_Weights.DEFAULT)
+    model.fc = nn.Linear(model.fc.in_features, NUM_CLASSES)
+
+    loaded = torch.load(CKPT_PATH, map_location=DEVICE)
+    state_dict = _normalize_state_dict(loaded)
+    model.load_state_dict(state_dict, strict=True)
+    model.to(DEVICE)
+    model.eval()
+    return model
+
+
+MODEL = _load_model()
+
+
+def sample_indices(num_frames: int, clip_len: int = CLIP_LEN) -> np.ndarray:
+    if num_frames <= clip_len:
+        return np.linspace(0, max(num_frames - 1, 0), clip_len).astype(int)
+    start = np.random.randint(0, num_frames - clip_len + 1)
+    return np.arange(start, start + clip_len)
+
+
+def get_clip_tensor(video_reader: VideoReader, inds: np.ndarray) -> torch.Tensor:
+    frames = video_reader.get_batch(inds).asnumpy()  # (T, H, W, C)
+    transformed: List[torch.Tensor] = []
+    for i in range(frames.shape[0]):
+        frame = EVAL_FRAME_TRANSFORM(frames[i])
+        frame = (frame - KIN_MEAN) / KIN_STD
+        transformed.append(frame)
+
+    # (C, T, H, W)
+    return torch.stack(transformed, dim=1)
+
+
+@torch.no_grad()
+def predict_path(video_path: str, n_clips: int = 10) -> Dict[str, object]:
+    vr = VideoReader(video_path, ctx=cpu(0))
+    total_frames = len(vr)
+    if total_frames == 0:
+        raise ValueError("Video has no readable frames.")
+
+    logits_sum = torch.zeros((1, NUM_CLASSES), device=DEVICE)
+    for _ in range(n_clips):
+        inds = sample_indices(total_frames, CLIP_LEN)
+        clip = get_clip_tensor(vr, inds).unsqueeze(0).to(DEVICE)
+        logits_sum += MODEL(clip)
+
+    logits = logits_sum / n_clips
+    probs = torch.softmax(logits, dim=1)[0].detach().cpu().numpy().tolist()
+    pred = int(np.argmax(probs))
+
+    return {
+        "prediction": "Stress (1)" if pred == 1 else "Normal (0)",
+        "probs": {"normal": probs[0], "stress": probs[1]},
+        "n_clips": n_clips,
+    }
+
+
+@app.get("/")
+def health() -> Dict[str, str]:
+    return {"status": "ok"}
+
+
+@app.post("/predict")
+async def predict_api(file: UploadFile = File(...)) -> Dict[str, object]:
+    suffix = os.path.splitext(file.filename or "")[1] or ".mp4"
+
+    with NamedTemporaryFile(delete=False, suffix=suffix) as tmp:
+        tmp.write(await file.read())
+        tmp_path = tmp.name
+
+    try:
+        return predict_path(tmp_path)
+    except Exception as exc:
+        raise HTTPException(status_code=400, detail=str(exc)) from exc
+    finally:
+        try:
+            os.remove(tmp_path)
+        except OSError:
+            pass

requirements.txt

@@ -0,0 +1,8 @@
+fastapi
+uvicorn[standard]
+python-multipart
+torch
+torchvision
+numpy
+pillow
+decord