import argparse import os import shutil import tempfile import cv2 from typing import Optional import numpy as np import pandas as pd import torch import torch.nn.functional as F from torch.utils import data from configs import LstmConfig, TransformerConfig from generate_keypoints import process_video from models import LSTM, Transformer from utils import load_json, load_label_map from video_preprocess import PreprocessConfig class KeypointsDataset(data.Dataset): def __init__(self, keypoints_dir, max_frame_len=169, frame_length=1080, frame_width=1920): self.files = sorted( [f for f in os.listdir(keypoints_dir) if f.endswith(".json")] ) self.keypoints_dir = keypoints_dir self.max_frame_len = max_frame_len self.frame_length = frame_length self.frame_width = frame_width def interpolate(self, arr): arr_x = arr[:, :, 0] arr_x = pd.DataFrame(arr_x) arr_x = arr_x.interpolate(method="linear", limit_direction="both").to_numpy() arr_y = arr[:, :, 1] arr_y = pd.DataFrame(arr_y) arr_y = arr_y.interpolate(method="linear", limit_direction="both").to_numpy() if np.count_nonzero(~np.isnan(arr_x)) == 0: arr_x = np.zeros(arr_x.shape) if np.count_nonzero(~np.isnan(arr_y)) == 0: arr_y = np.zeros(arr_y.shape) arr_x = arr_x * self.frame_width arr_y = arr_y * self.frame_length return np.stack([arr_x, arr_y], axis=-1) def combine_xy(self, x, y): x, y = np.array(x), np.array(y) _, length = x.shape x = x.reshape((-1, length, 1)) y = y.reshape((-1, length, 1)) return np.concatenate((x, y), -1).astype(np.float32) def __getitem__(self, idx): file_path = os.path.join(self.keypoints_dir, self.files[idx]) row = pd.read_json(file_path, typ="series") pose = self.combine_xy(row.pose_x, row.pose_y) h1 = self.combine_xy(row.hand1_x, row.hand1_y) h2 = self.combine_xy(row.hand2_x, row.hand2_y) pose = self.interpolate(pose) h1 = self.interpolate(h1) h2 = self.interpolate(h2) pose = pose.reshape(-1, 50).astype(np.float32) h1 = h1.reshape(-1, 42).astype(np.float32) h2 = h2.reshape(-1, 42).astype(np.float32) final_data = np.concatenate((pose, h1, h2), -1) final_data = np.pad( final_data, ((0, self.max_frame_len - final_data.shape[0]), (0, 0)), "constant", ) return {"uid": row.uid, "data": torch.FloatTensor(final_data)} def __len__(self): return len(self.files) def _pretrained_name(dataset: str, model_type: str, transformer_size: str) -> str: name = dataset name += "_no_cnn" if model_type == "lstm": name += "_lstm.pth" elif model_type == "transformer": name += "_transformer" name += "_large.pth" if transformer_size == "large" else "_small.pth" return name def load_model( dataset: str, model_type: str, transformer_size: str, checkpoint_path: Optional[str], label_map_path: Optional[str] = None, ): label_map = load_json(label_map_path) if label_map_path else load_label_map(dataset) n_classes = len(label_map) device = torch.device("cuda" if torch.cuda.is_available() else "cpu") if model_type == "transformer": config = TransformerConfig(size=transformer_size, max_position_embeddings=256) model = Transformer(config=config, n_classes=n_classes) elif model_type == "lstm": config = LstmConfig() model = LSTM(config=config, n_classes=n_classes) else: raise ValueError(f"Unsupported model_type: {model_type}") model = model.to(device) if checkpoint_path is None: pretrained_links = load_json("pretrained_links.json") model_name = _pretrained_name(dataset, model_type, transformer_size) # 1. First, check if you manually uploaded the file to the app folder if os.path.isfile(model_name): checkpoint_path = model_name else: # 2. If not, set the path to the writable /tmp directory checkpoint_path = os.path.join("/tmp", model_name) # 3. Download it to /tmp if it isn't already there if not os.path.isfile(checkpoint_path): link = pretrained_links.get(model_name) if not link or link == "link": raise FileNotFoundError(f"No pretrained link for {model_name}") torch.hub.download_url_to_file(link, checkpoint_path, progress=True) checkpoint_path = model_name ckpt = torch.load(checkpoint_path, map_location=device, weights_only=False) model.load_state_dict(ckpt["model"]) model.eval() return model, label_map @torch.no_grad() def predict_video( video_path: str, dataset: str, model, label_map: dict, preprocess_config: PreprocessConfig, *, max_frame_len: int = 169, use_holistic: bool = True, face_mode: str = "full", ): temp_dir = tempfile.mkdtemp(prefix="keypoints_") try: process_video( video_path, temp_dir, use_holistic=use_holistic, face_mode=face_mode, preprocess_config=preprocess_config, ) dataset_obj = KeypointsDataset(keypoints_dir=temp_dir, max_frame_len=max_frame_len) dataloader = data.DataLoader(dataset_obj, batch_size=1, shuffle=False) device = torch.device("cuda" if torch.cuda.is_available() else "cpu") id_to_label = dict(zip(label_map.values(), label_map.keys())) for batch in dataloader: input_data = batch["data"].to(device) logits = model(input_data) probs = torch.softmax(logits, dim=-1).detach().cpu().numpy()[0] idx = int(np.argmax(probs)) return { "uid": batch["uid"][0], "label": id_to_label[idx], "score": float(probs[idx]), } finally: shutil.rmtree(temp_dir, ignore_errors=True) raise RuntimeError("No frames processed from the video.") # smart brighten function that checks if the frame is too dark before applying CLAHE def smart_brighten(frame, darkness_threshold=80): """ Checks if a frame is too dark. If it is, applies CLAHE to enhance the hand/pose visibility without blowing out the highlights. """ # 1. Convert to grayscale just to calculate the average brightness gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY) average_brightness = np.mean(gray) # 2. If it's bright enough, leave it alone to save processing time if average_brightness > darkness_threshold: return frame # 3. If it's dark, apply CLAHE to the L channel (Lightness) # Convert BGR to LAB color space lab = cv2.cvtColor(frame, cv2.COLOR_BGR2LAB) l_channel, a, b = cv2.split(lab) # Apply CLAHE to the lightness channel clahe = cv2.createCLAHE(clipLimit=3.0, tileGridSize=(8,8)) cl = clahe.apply(l_channel) # Merge back and convert to BGR merged = cv2.merge((cl, a, b)) brightened_frame = cv2.cvtColor(merged, cv2.COLOR_LAB2BGR) return brightened_frame if __name__ == "__main__": parser = argparse.ArgumentParser(description="Run inference on a single video") parser.add_argument("--video", required=True, help="path to the video file") parser.add_argument( "--dataset", default="isl_split_dataset", help="dataset key used to resolve label map and pretrained model name", ) parser.add_argument("--model", default="transformer", choices=["transformer", "lstm"]) parser.add_argument("--transformer_size", default="small", choices=["small", "large"]) parser.add_argument("--checkpoint", default=None, help="path to model checkpoint") parser.add_argument( "--label_map_path", default=None, help="optional explicit label map JSON path", ) parser.add_argument("--apply_darken", action="store_true", help="apply darken before brighten") parser.add_argument("--darken_min", default=0.3, type=float) parser.add_argument("--darken_max", default=0.8, type=float) parser.add_argument("--brighten_method", default="clahe", choices=["clahe", "gamma"]) parser.add_argument("--brighten_gamma_min", default=1.2, type=float) parser.add_argument("--brighten_gamma_max", default=1.8, type=float) parser.add_argument( "--max_frame_len", default=169, type=int, help="sequence length used during training (must match inference)", ) args = parser.parse_args() preprocess_config = PreprocessConfig( apply_darken=args.apply_darken, apply_brighten=True, darken_min=args.darken_min, darken_max=args.darken_max, brighten_method=args.brighten_method, brighten_gamma_min=args.brighten_gamma_min, brighten_gamma_max=args.brighten_gamma_max, ) model, label_map = load_model( dataset=args.dataset, model_type=args.model, transformer_size=args.transformer_size, checkpoint_path=args.checkpoint, label_map_path=args.label_map_path, ) result = predict_video( args.video, args.dataset, model, label_map, preprocess_config=preprocess_config, max_frame_len=args.max_frame_len, ) print(result)