(image: FloatImage)
| 35 | |
| 36 | |
| 37 | def prep_image(image: FloatImage) -> UInt8[np.ndarray, "height width channel"]: |
| 38 | # Handle batched images. |
| 39 | if image.ndim == 4: |
| 40 | image = rearrange(image, "b c h w -> c h (b w)") |
| 41 | |
| 42 | # Handle single-channel images. |
| 43 | if image.ndim == 2: |
| 44 | image = rearrange(image, "h w -> () h w") |
| 45 | |
| 46 | # Ensure that there are 3 or 4 channels. |
| 47 | channel, _, _ = image.shape |
| 48 | if channel == 1: |
| 49 | image = repeat(image, "() h w -> c h w", c=3) |
| 50 | assert image.shape[0] in (3, 4) |
| 51 | |
| 52 | image = (image.detach().clip(min=0, max=1) * 255).type(torch.uint8) |
| 53 | return rearrange(image, "c h w -> h w c").cpu().numpy() |
| 54 | |
| 55 | |
| 56 | def save_image( |
no outgoing calls
no test coverage detected