| 120 | |
| 121 | |
| 122 | def stack_images(real, real_gens, gen, real_imgs=None): |
| 123 | # change to 3 channel |
| 124 | # print(real.shape) |
| 125 | # print(real_gens.shape) |
| 126 | # print(real_gens.shape) |
| 127 | # real = real[:3] |
| 128 | # real_gens = real_gens[:3] |
| 129 | # gen = gen[:3] |
| 130 | |
| 131 | nleft_cols = len(real_gens) + 1 |
| 132 | print("Stacking frames..") |
| 133 | allframes = np.concatenate( |
| 134 | (real[:, None, ...], *[x[:, None, ...] for x in real_gens], gen), 1) |
| 135 | nframes, nspa, nats, h, w, pix = allframes.shape |
| 136 | |
| 137 | blackborder = np.zeros((w // 30, h * nats, pix), dtype=allframes.dtype) |
| 138 | # blackborder = np.ones((w//30, h*nats, pix), dtype=allframes.dtype)*255 |
| 139 | frames = [] |
| 140 | for frame_idx in tqdm(range(nframes)): |
| 141 | columns = np.vstack(allframes[frame_idx].transpose(1, 2, 3, 4, |
| 142 | 0)).transpose( |
| 143 | 3, 1, 0, 2) |
| 144 | frame = np.concatenate( |
| 145 | (*columns[0:nleft_cols], blackborder, *columns[nleft_cols:]), |
| 146 | 0).transpose(1, 0, 2) |
| 147 | |
| 148 | frames.append(frame) |
| 149 | |
| 150 | if real_imgs is not None: |
| 151 | resize_imgs = convert_img(real_imgs, h)[:nframes, ...] |
| 152 | |
| 153 | for i in range(len(frames)): |
| 154 | imgs = np.vstack(resize_imgs[i, ...]) |
| 155 | imgs4 = np.ones( |
| 156 | (imgs.shape[0], imgs.shape[1], 4), dtype=np.uint8) * 255 |
| 157 | imgs4[:, :, :3] = imgs |
| 158 | #imgs = torch2numpy(imgs) |
| 159 | frames[i] = np.concatenate((imgs4, frames[i]), 1) |
| 160 | return np.stack(frames) |
| 161 | |
| 162 | |
| 163 | def stack_images_gen(gen, real_imgs=None): |