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.
(filelist: List[str], PIXEL_LIMIT: int = 255000, new_width: Optional[int] = None, verbose: bool = False)
| 16 | |
| 17 | |
| 18 | def load_images(filelist: List[str], PIXEL_LIMIT: int = 255000, new_width: Optional[int] = None, verbose: bool = False): |
| 19 | """ |
| 20 | Loads images from a directory or video, resizes them to a uniform size, |
| 21 | then converts and stacks them into a single [N, 3, H, W] PyTorch tensor. |
| 22 | """ |
| 23 | sources = [] |
| 24 | |
| 25 | # --- 1. Load image paths or video frames --- |
| 26 | for img_path in filelist: |
| 27 | try: |
| 28 | sources.append(Image.open(img_path).convert('RGB')) |
| 29 | except Exception as e: |
| 30 | print(f"Could not load image {img_path}: {e}") |
| 31 | |
| 32 | if not sources: |
| 33 | print("No images found or loaded.") |
| 34 | return torch.empty(0) |
| 35 | |
| 36 | if verbose: |
| 37 | print(f"Found {len(sources)} images/frames. Processing...") |
| 38 | |
| 39 | # --- 2. Determine a uniform target size for all images based on the first image --- |
| 40 | # This is necessary to ensure all tensors have the same dimensions for stacking. |
| 41 | first_img = sources[0] |
| 42 | W_orig, H_orig = first_img.size |
| 43 | if new_width is None: |
| 44 | scale = math.sqrt(PIXEL_LIMIT / (W_orig * H_orig)) if W_orig * H_orig > 0 else 1 |
| 45 | W_target, H_target = W_orig * scale, H_orig * scale |
| 46 | k, m = round(W_target / 14), round(H_target / 14) |
| 47 | while (k * 14) * (m * 14) > PIXEL_LIMIT: |
| 48 | if k / m > W_target / H_target: k -= 1 |
| 49 | else: m -= 1 |
| 50 | TARGET_W, TARGET_H = max(1, k) * 14, max(1, m) * 14 |
| 51 | else: |
| 52 | TARGET_W, TARGET_H = new_width, round(H_orig * (new_width / W_orig) / 14) * 14 |
| 53 | if verbose: |
| 54 | print(f"All images will be resized to a uniform size: ({TARGET_W}, {TARGET_H})") |
| 55 | |
| 56 | # --- 3. Resize images and convert them to tensors in the [0, 1] range --- |
| 57 | tensor_list = [] |
| 58 | # Define a transform to convert a PIL Image to a CxHxW tensor and normalize to [0,1] |
| 59 | to_tensor_transform = tvf.ToTensor() |
| 60 | |
| 61 | for img_pil in sources: |
| 62 | try: |
| 63 | # Resize to the uniform target size |
| 64 | resized_img = img_pil.resize((TARGET_W, TARGET_H), Image.Resampling.LANCZOS) |
| 65 | # Convert to tensor |
| 66 | img_tensor = to_tensor_transform(resized_img) |
| 67 | tensor_list.append(img_tensor) |
| 68 | except Exception as e: |
| 69 | print(f"Error processing an image: {e}") |
| 70 | |
| 71 | if not tensor_list: |
| 72 | print("No images were successfully processed.") |
| 73 | return torch.empty(0) |
| 74 | |
| 75 | # --- 4. Stack the list of tensors into a single [N, C, H, W] batch tensor --- |
no test coverage detected