import torch import torch.nn as nn from torchvision import datasets, transforms from torch.utils.data import Dataset, DataLoader import cv2 import numpy as np import os from pathlib import Path import warnings warnings.filterwarnings("ignore") class OptimizedVideoTransform: """Optimized video transforms for faster preprocessing""" def __init__(self, size=(112, 112), mean=None, std=None): self.size = size self.mean = mean or [0.43216, 0.394666, 0.37645] self.std = std or [0.22803, 0.22145, 0.216989] def __call__(self, clip): """Transform video clip to tensor format [C, T, H, W]""" if clip is None or len(clip) == 0: return torch.randn(3, 16, *self.size) # Convert list of frames to numpy array if needed if isinstance(clip, list): clip = np.stack(clip, axis=0) # Ensure clip has shape [T, H, W, C] if clip.ndim != 4 or clip.shape[-1] != 3: return torch.randn(3, 16, *self.size) # Resize frames efficiently using OpenCV resized_frames = [] for frame in clip: frame_resized = cv2.resize(frame, self.size, interpolation=cv2.INTER_AREA) resized_frames.append(frame_resized) # Convert to tensor and normalize clip_tensor = torch.from_numpy(np.stack(resized_frames)).float() / 255.0 # Permute from [T, H, W, C] to [C, T, H, W] clip_tensor = clip_tensor.permute(3, 0, 1, 2) # Normalize each channel for i in range(3): clip_tensor[i] = (clip_tensor[i] - self.mean[i]) / self.std[i] return clip_tensor class FastUCF101Dataset(Dataset): """Optimized UCF101 dataset loader with reduced frames and efficient loading""" def __init__(self, data_root, annotation_path, frames_per_clip=8, train=True, transform=None): self.data_root = Path(data_root) self.annotation_path = Path(annotation_path) self.frames_per_clip = frames_per_clip self.train = train self.transform = transform or OptimizedVideoTransform() # Load video paths and labels self.video_paths, self.labels = self._load_annotations() # Create class name mapping self.class_names = sorted(list(set([path.parent.name for path in self.video_paths]))) self.class_to_idx = {cls: idx for idx, cls in enumerate(self.class_names)} print(f"Loaded {len(self.video_paths)} {'train' if train else 'test'} videos") print(f"Classes: {len(self.class_names)}") def _load_annotations(self): """Load video paths and labels from annotation files""" if self.train: annotation_file = self.annotation_path / "trainlist01.txt" else: annotation_file = self.annotation_path / "testlist01.txt" video_paths = [] labels = [] with open(annotation_file, 'r') as f: for line in f: line = line.strip() if not line: continue if self.train: # Train file format: "classname/videoname.avi class_index" parts = line.split() video_path = parts[0] class_idx = int(parts[1]) - 1 # Convert to 0-based indexing else: # Test file format: "classname/videoname.avi" video_path = line class_name = video_path.split('/')[0] class_idx = self.class_names.index(class_name) if hasattr(self, 'class_names') else 0 full_path = self.data_root / video_path if full_path.exists(): video_paths.append(full_path) labels.append(class_idx) return video_paths, labels def _load_video_frames(self, video_path): """Efficiently load video frames using OpenCV""" cap = cv2.VideoCapture(str(video_path)) frames = [] # Get total frame count total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) if total_frames == 0: cap.release() return None # Sample frames uniformly if total_frames >= self.frames_per_clip: frame_indices = np.linspace(0, total_frames - 1, self.frames_per_clip, dtype=int) else: # Repeat frames if video is too short frame_indices = np.arange(total_frames) while len(frame_indices) < self.frames_per_clip: frame_indices = np.concatenate([frame_indices, frame_indices[:self.frames_per_clip - len(frame_indices)]]) for i, frame_idx in enumerate(frame_indices): cap.set(cv2.CAP_PROP_POS_FRAMES, frame_idx) ret, frame = cap.read() if ret: # Convert BGR to RGB frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) frames.append(frame) else: # If frame reading fails, duplicate the last frame if frames: frames.append(frames[-1]) else: # Create a black frame as fallback frames.append(np.zeros((240, 320, 3), dtype=np.uint8)) cap.release() return frames[:self.frames_per_clip] def __len__(self): return len(self.video_paths) def __getitem__(self, idx): try: video_path = self.video_paths[idx] label = self.labels[idx] # Load video frames frames = self._load_video_frames(video_path) if frames is None: # Return dummy data on error return self.transform(None), 0 # Apply transforms video_tensor = self.transform(frames) return video_tensor, label except Exception as e: if idx < 5: # Only print first few errors print(f"Error loading video {idx}: {e}") return self.transform(None), 0 def get_optimized_datasets(data_root, annotation_path, frames_per_clip=8): """Get optimized train and test datasets""" transform = OptimizedVideoTransform(size=(112, 112)) print("Loading optimized training dataset...") train_dataset = FastUCF101Dataset( data_root, annotation_path, frames_per_clip=frames_per_clip, train=True, transform=transform ) print("Loading optimized test dataset...") test_dataset = FastUCF101Dataset( data_root, annotation_path, frames_per_clip=frames_per_clip, train=False, transform=transform ) return train_dataset, test_dataset def get_optimized_dataloaders(data_root, annotation_path, batch_size=16, frames_per_clip=8, num_workers=4): """Get optimized dataloaders with better performance settings""" train_dataset, test_dataset = get_optimized_datasets(data_root, annotation_path, frames_per_clip) # Configure DataLoader parameters based on num_workers if num_workers > 0: # Multi-processing configuration train_kwargs = { 'persistent_workers': True, 'prefetch_factor': 2 } test_kwargs = { 'persistent_workers': True, 'prefetch_factor': 2 } else: # Single-process configuration (Windows safe) train_kwargs = {} test_kwargs = {} # Create DataLoaders with correct configuration train_loader = DataLoader( train_dataset, batch_size=batch_size, shuffle=True, num_workers=num_workers, pin_memory=True if torch.cuda.is_available() else False, **train_kwargs ) test_loader = DataLoader( test_dataset, batch_size=batch_size, shuffle=False, num_workers=num_workers, pin_memory=True if torch.cuda.is_available() else False, **test_kwargs ) return train_loader, test_loader, train_dataset.class_names