Loads images from a directory or video, resizes them to a uniform size, then converts and stacks them into a single [N, 3, H, W] PyTorch tensor.
(path='data/truck', interval=1, PIXEL_LIMIT=255000)
| 249 | |
| 250 | |
| 251 | def load_images_as_tensor(path='data/truck', interval=1, PIXEL_LIMIT=255000): |
| 252 | """ |
| 253 | Loads images from a directory or video, resizes them to a uniform size, |
| 254 | then converts and stacks them into a single [N, 3, H, W] PyTorch tensor. |
| 255 | """ |
| 256 | sources = [] |
| 257 | |
| 258 | # --- 1. Load image paths or video frames --- |
| 259 | if osp.isdir(path): |
| 260 | print(f"Loading images from directory: {path}") |
| 261 | filenames = sorted([x for x in os.listdir(path) if x.lower().endswith(('.png', '.jpg', '.jpeg'))]) |
| 262 | for i in range(0, len(filenames), interval): |
| 263 | img_path = osp.join(path, filenames[i]) |
| 264 | try: |
| 265 | sources.append(Image.open(img_path).convert('RGB')) |
| 266 | except Exception as e: |
| 267 | print(f"Could not load image {filenames[i]}: {e}") |
| 268 | elif path.lower().endswith('.mp4'): |
| 269 | print(f"Loading frames from video: {path}") |
| 270 | cap = cv2.VideoCapture(path) |
| 271 | if not cap.isOpened(): raise IOError(f"Cannot open video file: {path}") |
| 272 | frame_idx = 0 |
| 273 | while True: |
| 274 | ret, frame = cap.read() |
| 275 | if not ret: break |
| 276 | if frame_idx % interval == 0: |
| 277 | rgb_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) |
| 278 | sources.append(Image.fromarray(rgb_frame)) |
| 279 | frame_idx += 1 |
| 280 | cap.release() |
| 281 | else: |
| 282 | raise ValueError(f"Unsupported path. Must be a directory or a .mp4 file: {path}") |
| 283 | |
| 284 | if not sources: |
| 285 | print("No images found or loaded.") |
| 286 | return torch.empty(0) |
| 287 | |
| 288 | print(f"Found {len(sources)} images/frames. Processing...") |
| 289 | |
| 290 | # --- 2. Determine a uniform target size for all images based on the first image --- |
| 291 | # This is necessary to ensure all tensors have the same dimensions for stacking. |
| 292 | first_img = sources[0] |
| 293 | W_orig, H_orig = first_img.size |
| 294 | scale = math.sqrt(PIXEL_LIMIT / (W_orig * H_orig)) if W_orig * H_orig > 0 else 1 |
| 295 | W_target, H_target = W_orig * scale, H_orig * scale |
| 296 | k, m = round(W_target / 14), round(H_target / 14) |
| 297 | while (k * 14) * (m * 14) > PIXEL_LIMIT: |
| 298 | if k / m > W_target / H_target: k -= 1 |
| 299 | else: m -= 1 |
| 300 | TARGET_W, TARGET_H = max(1, k) * 14, max(1, m) * 14 |
| 301 | print(f"All images will be resized to a uniform size: ({TARGET_W}, {TARGET_H})") |
| 302 | |
| 303 | # --- 3. Resize images and convert them to tensors in the [0, 1] range --- |
| 304 | tensor_list = [] |
| 305 | # Define a transform to convert a PIL Image to a CxHxW tensor and normalize to [0,1] |
| 306 | to_tensor_transform = transforms.ToTensor() |
| 307 | |
| 308 | for img_pil in sources: |