import os import sys import time import torch import numpy as np from dotenv import load_dotenv if hasattr(sys.stdout, 'reconfigure'): try: sys.stdout.reconfigure(encoding='utf-8', errors='replace') except Exception: pass def format_log(text: str, start_time: float = None, target_col: int = 64) -> str: """Formats a log message with tabulators aligning timestamps in a clean column.""" if start_time is not None: elapsed = time.time() - start_time time_str = f"[{elapsed:6.2f} s since start]" else: time_str = "" num_tabs = max(1, (target_col - len(text) + 7) // 8) tabs = "\t" * num_tabs return f"{text}{tabs}{time_str}" # Import the translation model and helper functions from Uni_Sign from Uni_Sign.models import Uni_Sign from Uni_Sign.datasets import load_part_kp_YTASL, YTASL_GROUP_SIZES, _fill_missing_landmarks, select_frame_indices from predict_pose import create_mediapipe_models, predict_pose, load_video_cv from visualize_pose import render_keypoints_video os.environ['KMP_DUPLICATE_LIB_OK'] = 'TRUE' load_dotenv() # Global variables to cache models and config so they are not reloaded on every video translation model = None pose_models = None args = None class InferenceConfig: """Configuration class for Uni_Sign model inference.""" def __init__(self): # Path to the pre-trained weights checkpoint self.finetune = os.environ.get("UNISIGN_WEIGHTS", r"./Uni_Sign/unisign_model/best_checkpoint-wlasl.pth") self.dataset = "YTASL" self.task = "ISLR" self.max_length = 256 self.normalization = "none" self.layout = "pruned" self.n_registers = 0 self.hidden_dim = 256 self.rgb_support = False self.no_adaptive_gcn = False self.register_position = "before_all" self.label_smoothing = 0.2 self.batch_size = 1 def initialize_model(): """Initializes the Uni_Sign translator and the Mediapipe/YOLO pose estimation models.""" global model, pose_models, args if model is not None and pose_models is not None: return print("Initializing Uni_Sign model...") args = InferenceConfig() model = Uni_Sign(args=args) # Load model weights if args.finetune and os.path.exists(args.finetune): print(f"Loading checkpoint weights from: {args.finetune}") state_dict = torch.load(args.finetune, map_location='cpu')['model'] model.load_state_dict(state_dict, strict=False) else: print(f"WARNING: Checkpoint weights not found at '{args.finetune}'!") device = "cuda" if torch.cuda.is_available() else "cpu" model.to(device) model.eval() print("Initializing Mediapipe and YOLO models...") pose_checkpoint_folder = 'checkpoints/pose/' pose_models = create_mediapipe_models(pose_checkpoint_folder) def process_pose_data_in_memory(pose_results, args): """ Processes pose keypoints in-memory from dictionary results. Converts keypoint data to Tensors formatted for Uni_Sign input. """ raw_pose = pose_results.get('cropped_keypoints', []) if not raw_pose: raise ValueError("Video does not contain any detected pose (cropped_keypoints).") # Format raw coordinates into standard [X, Y] coordinates for each body part pose = [] for frame_data in raw_pose: formatted_frame = {} for part, expected_size in YTASL_GROUP_SIZES.items(): kps = frame_data.get(part, []) # Use empty list if data is missing; _fill_missing_landmarks will fill it if kps is None or len(kps) == 0: formatted_frame[part] = [] else: # Standardize to only X and Y coordinates (ignoring confidence/Z columns) formatted_frame[part] = np.array(kps)[:, :2].tolist() pose.append(formatted_frame) # Subsample frames to match target length (max_length) duration = len(pose) tmp = select_frame_indices(duration, args.max_length, phase='test') skeletons = [pose[i] for i in tmp] # Fill in missing landmarks to maintain consistent shape dimensions confs = [] for i, skeleton in enumerate(skeletons): conf = {} for group_name, expected_size in YTASL_GROUP_SIZES.items(): _fill_missing_landmarks( skeleton=skeleton, conf=conf, group_name=group_name, expected_size=expected_size, clip_name="gradio_video", frame_idx=i, error_group_label=f"group '{group_name}'", include_size_details=False, strict_key_access=True, ) confs.append(conf) # Normalize keypoints and structure them for YTASL dataset kps_with_scores = load_part_kp_YTASL(skeletons, confs, args.normalization, args.layout) # Batch features by unsqueezing to add batch dimension (batch_size = 1) src_input = {} for key, val in kps_with_scores.items(): src_input[key] = val.unsqueeze(0) # Generate attention mask and sequence lengths seq_len = src_input['body'].shape[1] mask_gen = torch.ones([seq_len]) + 7 src_input['attention_mask'] = (mask_gen != 0).long().unsqueeze(0) src_input['src_length_batch'] = torch.LongTensor([seq_len]) src_input['name_batch'] = ["gradio_video"] return src_input def process_input(input_video_path, progress=None, log_callback=None): """ Main entry point for Gradio. Runs pose extraction, pre-processing, and model translation inference via Uni_Sign. """ try: start_time = time.time() def emit(text, prog_val=None, step_title=None): formatted = format_log(text, start_time) print(formatted) if progress is not None and prog_val is not None: try: progress(prog_val, desc=step_title or text) except Exception: pass if log_callback is not None: try: log_callback(formatted, prog_val, step_title or text) except Exception: pass emit("0. Initializing models...", 0.02, "0. Initializing models...") initialize_model() # Check if precomputed keypoints JSON & pre-rendered video exist (e.g. from examples.json or file pattern) video_dir = os.path.dirname(os.path.abspath(input_video_path)) video_stem = os.path.splitext(os.path.basename(input_video_path))[0] # Normalize stem in case it ends with _keypoints clean_stem = video_stem.replace("_keypoints", "") possible_json = os.path.join(video_dir, f"{clean_stem}_keypoints.json") possible_video = os.path.join(video_dir, f"{clean_stem}_keypoints.mp4") # Also check examples.json registry examples_json_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "examples.json") if os.path.exists(examples_json_path): try: import json with open(examples_json_path, "r", encoding="utf-8") as f: ex_data = json.load(f) for item in ex_data: if os.path.normpath(item.get("video_path", "")) == os.path.normpath(input_video_path): if item.get("keypoints_path") and os.path.exists(item["keypoints_path"]): possible_json = item["keypoints_path"] if item.get("keypoints_video_path") and os.path.exists(item["keypoints_video_path"]): possible_video = item["keypoints_video_path"] break except Exception: pass if os.path.exists(possible_json) and os.path.getsize(possible_json) > 0: emit("1. Loading precomputed keypoints from JSON...", 0.50, "1. Loading precomputed keypoints...") import json with open(possible_json, "r", encoding="utf-8") as f: pose_results = json.load(f) keypoints_video_path = possible_video if os.path.exists(possible_video) else "" emit(" [1e] Precomputed keypoints loaded from cache", 0.90, " [1e] Precomputed keypoints loaded") else: emit("1. Extracting keypoints from video...", 0.05, "1. Extracting keypoints from video...") video_frames, fps = load_video_cv(input_video_path) if fps is None or fps <= 0 or np.isnan(fps): fps = 25.0 num_frames = len(video_frames) dur = num_frames / fps if fps > 0 else 0 emit(f" [1a] Video loaded: {num_frames} frames ({fps:.1f} FPS, duration {dur:.1f} s)", 0.06, f" [1a] Video loaded ({num_frames} frames)") pose_results = predict_pose(video_frames, pose_models, progress=progress, start_time=start_time, log_callback=log_callback) # Attach bounding boxes to each frame's keypoints kps_list = pose_results.get("keypoints", []) bbox_face_list = pose_results.get("bbox_face", []) bbox_lh_list = pose_results.get("bbox_left_hand", []) bbox_rh_list = pose_results.get("bbox_right_hand", []) for i, kp in enumerate(kps_list): if i < len(bbox_face_list): kp['bbox_face'] = bbox_face_list[i] if i < len(bbox_lh_list): kp['bbox_left_hand'] = bbox_lh_list[i] if i < len(bbox_rh_list): kp['bbox_right_hand'] = bbox_rh_list[i] # Render keypoint overlay video keypoints_video_path = "" try: keypoints_video_path = render_keypoints_video(video_frames, kps_list, fps=fps, progress=progress, start_time=start_time, log_callback=log_callback) emit(" [1e] Keypoint video saved successfully", 0.90, " [1e] Keypoint video saved") except Exception as kp_err: emit(f" [1e] Warning rendering keypoint video: {kp_err}", 0.90, " [1e] Warning keypoint video") emit("2. Pre-processing visual features...", 0.92, "2. Pre-processing visual features...") emit("3. Converting keypoints to Tensors for Uni-Sign...", 0.94, "3. Converting keypoints to Tensors...") src_input = process_pose_data_in_memory(pose_results, args) # Move Tensors to the same device (GPU/CPU) as the model device = next(model.parameters()).device for key in ['body', 'left', 'right', 'face_all', 'attention_mask']: if key in src_input: # Cast features to float32 (except the attention mask) if key != 'attention_mask': src_input[key] = src_input[key].to(device, dtype=torch.float32) else: src_input[key] = src_input[key].to(device) tgt_input = {'gt_sentence': [""], 'gt_gloss': [""]} emit("4. Generating translation with Uni-Sign...", 0.96, "4. Generating translation...") with torch.no_grad(): stack_out = model(src_input, tgt_input) output = model.generate( stack_out, max_new_tokens=100, num_beams=4, ) tokenizer = model.mt5_tokenizer tgt_pres = tokenizer.batch_decode(output, skip_special_tokens=True) result = tgt_pres[0].strip() if not result: result = "The model was unable to generate a translation. Please try a different video." emit("5. Translation completed!", 1.0, "5. Translation completed!") return result, keypoints_video_path except Exception as e: print(f"Error during video processing: {e}") import traceback traceback.print_exc() return f"Error processing video: {str(e)}", ""