open and convert all images in a list or folder to proper input format for DUSt3R
(folder_or_list, size, square_ok=False, verbose=True, patch_size=16)
| 72 | |
| 73 | |
| 74 | def load_images(folder_or_list, size, square_ok=False, verbose=True, patch_size=16): |
| 75 | """ open and convert all images in a list or folder to proper input format for DUSt3R |
| 76 | """ |
| 77 | if isinstance(folder_or_list, str): |
| 78 | if verbose: |
| 79 | print(f'>> Loading images from {folder_or_list}') |
| 80 | root, folder_content = folder_or_list, sorted(os.listdir(folder_or_list)) |
| 81 | |
| 82 | elif isinstance(folder_or_list, list): |
| 83 | if verbose: |
| 84 | print(f'>> Loading a list of {len(folder_or_list)} images') |
| 85 | root, folder_content = '', folder_or_list |
| 86 | |
| 87 | else: |
| 88 | raise ValueError(f'bad {folder_or_list=} ({type(folder_or_list)})') |
| 89 | |
| 90 | supported_images_extensions = ['.jpg', '.jpeg', '.png'] |
| 91 | if heif_support_enabled: |
| 92 | supported_images_extensions += ['.heic', '.heif'] |
| 93 | supported_images_extensions = tuple(supported_images_extensions) |
| 94 | |
| 95 | imgs = [] |
| 96 | for path in folder_content: |
| 97 | if not path.lower().endswith(supported_images_extensions): |
| 98 | continue |
| 99 | img = exif_transpose(PIL.Image.open(os.path.join(root, path))).convert('RGB') |
| 100 | W1, H1 = img.size |
| 101 | if size == 224: |
| 102 | # resize short side to 224 (then crop) |
| 103 | img = _resize_pil_image(img, round(size * max(W1/H1, H1/W1))) |
| 104 | else: |
| 105 | # resize long side to 512 |
| 106 | img = _resize_pil_image(img, size) |
| 107 | W, H = img.size |
| 108 | cx, cy = W//2, H//2 |
| 109 | if size == 224: |
| 110 | half = min(cx, cy) |
| 111 | img = img.crop((cx-half, cy-half, cx+half, cy+half)) |
| 112 | else: |
| 113 | halfw = ((2 * cx) // patch_size) * patch_size / 2 |
| 114 | halfh = ((2 * cy) // patch_size) * patch_size / 2 |
| 115 | if not (square_ok) and W == H: |
| 116 | halfh = 3*halfw/4 |
| 117 | img = img.crop((cx-halfw, cy-halfh, cx+halfw, cy+halfh)) |
| 118 | |
| 119 | W2, H2 = img.size |
| 120 | if verbose: |
| 121 | print(f' - adding {path} with resolution {W1}x{H1} --> {W2}x{H2}') |
| 122 | imgs.append(dict(img=ImgNorm(img)[None], true_shape=np.int32( |
| 123 | [img.size[::-1]]), idx=len(imgs), instance=str(len(imgs)))) |
| 124 | |
| 125 | assert imgs, 'no images foud at '+root |
| 126 | if verbose: |
| 127 | print(f' (Found {len(imgs)} images)') |
| 128 | return imgs |
no test coverage detected