Args: ims: Tensor of shape (b, c, h, w) stack: "row" or "col" split: If 'row' stack by rows, if 'col' stack by columns. Returns: Tensor of shape (h, w, c)
(ims, stack="row", split=4, channel_last=False)
| 5 | |
| 6 | |
| 7 | def ims_to_grid(ims, stack="row", split=4, channel_last=False): |
| 8 | """ |
| 9 | Args: |
| 10 | ims: Tensor of shape (b, c, h, w) |
| 11 | stack: "row" or "col" |
| 12 | split: If 'row' stack by rows, if 'col' stack by columns. |
| 13 | Returns: |
| 14 | Tensor of shape (h, w, c) |
| 15 | """ |
| 16 | if stack not in ["row", "col"]: |
| 17 | raise ValueError(f"Unknown stack type {stack}") |
| 18 | from_ = 'h w c' if channel_last else 'c h w' |
| 19 | if split is not None and ims.shape[0] % split == 0: |
| 20 | splitter = dict(b1=split) if stack == "row" else dict(b2=split) |
| 21 | ims = einops.rearrange(ims, f"(b1 b2) {from_} -> (b1 h) (b2 w) c", **splitter) |
| 22 | else: |
| 23 | to = "(b h) w c" if stack == "row" else "h (b w) c" |
| 24 | ims = einops.rearrange(ims, f"b {from_} -> " + to) |
| 25 | return ims |
| 26 | |
| 27 | |
| 28 | def tensor2im(tensor, denormalize_zero_one=False): |