Skylorjustine/Video-Action-Recognition

1#!/usr/bin/env python32"""3Fixed video action prediction with proper TimeSformer tensor format.4This version resolves the tensor compatibility issues definitively.5"""6 7import argparse8import json9import logging10from pathlib import Path11from typing import List, Tuple, Optional12import warnings13 14# Suppress warnings for cleaner output15warnings.filterwarnings("ignore", category=UserWarning)16warnings.filterwarnings("ignore", category=DeprecationWarning)17 18import torch19from PIL import Image20 21# Configure logging22logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')23 24# Video reading libraries25try:26    import cv227    HAS_CV2 = True28except ImportError:29    HAS_CV2 = False30    cv2 = None31 32try:33    import decord34    HAS_DECORD = True35except ImportError:36    HAS_DECORD = False37    decord = None38 39MODEL_ID = "facebook/timesformer-base-finetuned-k400"40 41def read_video_frames_cv2(video_path: Path, num_frames: int = 8) -> List[Image.Image]:42    """Read frames using OpenCV with robust error handling."""43    if not HAS_CV2:44        raise RuntimeError("OpenCV not available")45 46    cap = cv2.VideoCapture(str(video_path))47    if not cap.isOpened():48        raise RuntimeError(f"Cannot open video: {video_path}")49 50    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))51    if total_frames == 0:52        cap.release()53        raise RuntimeError("Video has no frames")54 55    # Sample frames uniformly across the video56    if total_frames <= num_frames:57        frame_indices = list(range(total_frames))58    else:59        step = max(1, total_frames // num_frames)60        frame_indices = [i * step for i in range(num_frames)]61        # Ensure we don't exceed total frames62        frame_indices = [min(idx, total_frames - 1) for idx in frame_indices]63 64    frames = []65    for idx in frame_indices:66        cap.set(cv2.CAP_PROP_POS_FRAMES, idx)67        ret, frame = cap.read()68        if ret:69            # Convert BGR to RGB70            frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)71            pil_image = Image.fromarray(frame_rgb)72            frames.append(pil_image)73 74    cap.release()75 76    # Pad with last frame if needed77    while len(frames) < num_frames:78        if frames:79            frames.append(frames[-1].copy())80        else:81            # Create black frame as fallback82            black_frame = Image.new('RGB', (224, 224), (0, 0, 0))83            frames.append(black_frame)84 85    return frames[:num_frames]86 87def read_video_frames_decord(video_path: Path, num_frames: int = 8) -> List[Image.Image]:88    """Read frames using decord."""89    if not HAS_DECORD:90        raise RuntimeError("Decord not available")91 92    vr = decord.VideoReader(str(video_path))93    total_frames = len(vr)94 95    if total_frames == 0:96        raise RuntimeError("Video has no frames")97 98    # Sample frames99    if total_frames <= num_frames:100        indices = list(range(total_frames))101    else:102        step = max(1, total_frames // num_frames)103        indices = [i * step for i in range(num_frames)]104        indices = [min(idx, total_frames - 1) for idx in indices]105 106    try:107        frame_arrays = vr.get_batch(indices).asnumpy()108        frames = [Image.fromarray(frame) for frame in frame_arrays]109    except Exception:110        # Fallback to individual frame reading111        frames = []112        for idx in indices:113            try:114                frame = vr[idx].asnumpy()115                frames.append(Image.fromarray(frame))116            except Exception:117                continue118 119    # Pad if necessary120    while len(frames) < num_frames:121        if frames:122            frames.append(frames[-1].copy())123        else:124            black_frame = Image.new('RGB', (224, 224), (0, 0, 0))125            frames.append(black_frame)126 127    return frames[:num_frames]128 129def read_video_frames(video_path: Path, num_frames: int = 8) -> List[Image.Image]:130    """Read video frames with fallback methods."""131    last_error = None132 133    # Try decord first (usually faster and more reliable)134    if HAS_DECORD:135        try:136            frames = read_video_frames_decord(video_path, num_frames)137            if frames and len(frames) > 0:138                logging.debug(f"Successfully read {len(frames)} frames using decord")139                return frames140        except Exception as e:141            last_error = e142            logging.debug(f"Decord failed: {e}")143 144    # Fallback to OpenCV145    if HAS_CV2:146        try:147            frames = read_video_frames_cv2(video_path, num_frames)148            if frames and len(frames) > 0:149                logging.debug(f"Successfully read {len(frames)} frames using OpenCV")150                return frames151        except Exception as e:152            last_error = e153            logging.debug(f"OpenCV failed: {e}")154 155    if last_error:156        raise RuntimeError(f"Failed to read video frames: {last_error}")157    else:158        raise RuntimeError("No video reading library available")159 160def normalize_frames(frames: List[Image.Image], target_size: Tuple[int, int] = (224, 224)) -> List[Image.Image]:161    """Normalize frames to consistent format."""162    if not frames:163        raise RuntimeError("No frames to normalize")164 165    normalized = []166    for i, frame in enumerate(frames):167        try:168            # Convert to RGB if needed169            if frame.mode != 'RGB':170                frame = frame.convert('RGB')171 172            # Resize to target size173            if frame.size != target_size:174                frame = frame.resize(target_size, Image.Resampling.LANCZOS)175 176            normalized.append(frame)177        except Exception as e:178            logging.warning(f"Error normalizing frame {i}: {e}")179            # Create a black frame as fallback180            black_frame = Image.new('RGB', target_size, (0, 0, 0))181            normalized.append(black_frame)182 183    return normalized184 185def create_timesformer_tensor(frames: List[Image.Image]) -> torch.Tensor:186    """187    Create properly formatted tensor for TimeSformer model.188 189    TimeSformer expects 5D input tensor:190    Input format: [batch_size, num_frames, channels, height, width]191    For 8 frames of 224x224: [1, 8, 3, 224, 224]192    """193    if len(frames) != 8:194        raise ValueError(f"Expected 8 frames, got {len(frames)}")195 196    # Convert frames to tensors without using numpy197    frame_tensors = []198 199    for frame in frames:200        # Ensure correct format201        if frame.mode != 'RGB':202            frame = frame.convert('RGB')203        if frame.size != (224, 224):204            frame = frame.resize((224, 224), Image.Resampling.LANCZOS)205 206        # Convert PIL image to tensor manually to avoid numpy issues207        pixels = list(frame.getdata())  # List of (R, G, B) tuples208 209        # Separate into RGB channels and normalize210        r_channel = []211        g_channel = []212        b_channel = []213 214        for r, g, b in pixels:215            r_channel.append(r / 255.0)216            g_channel.append(g / 255.0)217            b_channel.append(b / 255.0)218 219        # Reshape to 2D (224, 224) for each channel220        r_tensor = torch.tensor(r_channel, dtype=torch.float32).view(224, 224)221        g_tensor = torch.tensor(g_channel, dtype=torch.float32).view(224, 224)222        b_tensor = torch.tensor(b_channel, dtype=torch.float32).view(224, 224)223 224        # Stack channels: (3, 224, 224)225        frame_tensor = torch.stack([r_tensor, g_tensor, b_tensor], dim=0)226        frame_tensors.append(frame_tensor)227 228    # Stack frames: (8, 3, 224, 224)229    video_tensor = torch.stack(frame_tensors, dim=0)230 231    # Rearrange to TimeSformer format: (batch, frames, channels, height, width)232    # From (8, 3, 224, 224) to (1, 8, 3, 224, 224)233    video_tensor = video_tensor.unsqueeze(0)  # Add batch dimension: (1, 8, 3, 224, 224)234 235    logging.debug(f"Created tensor with shape: {video_tensor.shape}")236    logging.debug(f"Tensor dtype: {video_tensor.dtype}")237    logging.debug(f"Tensor range: [{video_tensor.min():.3f}, {video_tensor.max():.3f}]")238 239    return video_tensor240 241def load_model(device: Optional[str] = None):242    """Load TimeSformer model and processor."""243    try:244        from transformers import AutoImageProcessor, TimesformerForVideoClassification245 246        device = device or ("cuda" if torch.cuda.is_available() else "cpu")247        logging.info(f"Loading model on device: {device}")248 249        processor = AutoImageProcessor.from_pretrained(MODEL_ID)250        model = TimesformerForVideoClassification.from_pretrained(MODEL_ID)251        model.to(device)252        model.eval()253 254        logging.info("Model loaded successfully")255        return processor, model, device256 257    except Exception as e:258        logging.error(f"Failed to load model: {e}")259        raise RuntimeError(f"Model loading failed: {e}")260 261def predict_actions(video_path: str, top_k: int = 5) -> List[Tuple[str, float]]:262    """263    Predict actions in video using TimeSformer model.264 265    Args:266        video_path: Path to video file267        top_k: Number of top predictions to return268 269    Returns:270        List of (action_label, confidence_score) tuples271    """272    video_path = Path(video_path)273 274    if not video_path.exists():275        raise FileNotFoundError(f"Video file not found: {video_path}")276 277    try:278        # Load model279        processor, model, device = load_model()280 281        # Extract and normalize frames282        logging.info(f"Processing video: {video_path.name}")283        frames = read_video_frames(video_path, num_frames=8)284        frames = normalize_frames(frames, target_size=(224, 224))285 286        logging.info(f"Extracted and normalized {len(frames)} frames")287 288        # Create tensor in correct format289        pixel_values = create_timesformer_tensor(frames)290        pixel_values = pixel_values.to(device)291 292        # Run inference293        logging.info("Running model inference...")294        with torch.no_grad():295            outputs = model(pixel_values=pixel_values)296            logits = outputs.logits297 298        # Get top-k predictions299        probabilities = torch.softmax(logits, dim=-1)[0]  # Remove batch dimension300        top_probs, top_indices = torch.topk(probabilities, k=top_k)301 302        # Convert to results303        results = []304        for prob, idx in zip(top_probs, top_indices):305            label = model.config.id2label[idx.item()]306            confidence = float(prob.item())307            results.append((label, confidence))308 309        logging.info(f"Generated {len(results)} predictions successfully")310 311        # Log top prediction for debugging312        if results:313            top_label, top_conf = results[0]314            logging.info(f"Top prediction: {top_label} ({top_conf:.3f})")315 316        return results317 318    except Exception as e:319        logging.error(f"Prediction failed: {e}")320        raise RuntimeError(f"Video processing error: {e}")321 322def main():323    """Command line interface."""324    parser = argparse.ArgumentParser(description="Predict actions in video using TimeSformer")325    parser.add_argument("video", type=str, help="Path to video file")326    parser.add_argument("--top-k", type=int, default=5, help="Number of top predictions")327    parser.add_argument("--json", action="store_true", help="Output as JSON")328    parser.add_argument("--verbose", "-v", action="store_true", help="Enable verbose logging")329 330    args = parser.parse_args()331 332    if args.verbose:333        logging.getLogger().setLevel(logging.DEBUG)334 335    try:336        # Run prediction337        predictions = predict_actions(args.video, top_k=args.top_k)338 339        if args.json:340            output = [{"label": label, "confidence": confidence}341                     for label, confidence in predictions]342            print(json.dumps(output, indent=2))343        else:344            print(f"\nTop {len(predictions)} predictions for: {args.video}")345            print("-" * 60)346            for i, (label, confidence) in enumerate(predictions, 1):347                print(f"{i:2d}. {label:<35} {confidence:.4f}")348 349        return 0350 351    except Exception as e:352        print(f"Error: {e}")353        if args.verbose:354            import traceback355            traceback.print_exc()356        return 1357 358if __name__ == "__main__":359    exit(main())360