Spaces:
Sleeping
Sleeping
File size: 4,353 Bytes
20acb83 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 | 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
|