Spaces:
Sleeping
Sleeping
Add initial project files including Dockerfile, main application logic, model checkpoint, and requirements
20acb83 | 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) | |
| 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, | |
| } | |
| def health() -> Dict[str, str]: | |
| return {"status": "ok"} | |
| 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 | |