Source code for imaginaire.generators.gancraft

# Copyright (C) 2021 NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
# This work is made available under the Nvidia Source Code License-NC.
# To view a copy of this license, check out LICENSE.md
import os

import cv2
import imageio
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F

import imaginaire.model_utils.gancraft.camctl as camctl
import imaginaire.model_utils.gancraft.mc_utils as mc_utils
import imaginaire.model_utils.gancraft.voxlib as voxlib
from imaginaire.utils.distributed import master_only_print as print
from imaginaire.generators.gancraft_base import Base3DGenerator, RenderMLP  # noqa


[docs]class Generator(Base3DGenerator): r"""GANcraft generator constructor. Args: gen_cfg (obj): Generator definition part of the yaml config file. data_cfg (obj): Data definition part of the yaml config file. """ def __init__(self, gen_cfg, data_cfg): super(Generator, self).__init__(gen_cfg, data_cfg) print('GANcraft generator initialization.') # Load voxels of the input world. # The loaded voxel tensor has a shape of [X, Y, Z], dtype==torch.int32 # 0 means empty (air). print('[Generator] Loading voxel world: ', gen_cfg.voxel_path) if gen_cfg.voxel_path.endswith('.npy'): voxel_t = np.load(gen_cfg.voxel_path) voxel_t = torch.from_numpy(voxel_t.astype(np.int32)) else: voxel_t = mc_utils.load_voxel_new(gen_cfg.voxel_path, shape=gen_cfg.voxel_shape) print('[Generator] Loaded voxel world.') self.voxel = mc_utils.McVoxel(voxel_t, preproc_ver=gen_cfg.voxel_preproc_ver) blk_feats = torch.empty([self.voxel.nfilledvox, gen_cfg.blk_feat_dim], requires_grad=True) self.blk_feats = nn.Parameter(blk_feats) # Feature per voxel corner. # Minecraft -> SPADE label translator. self.label_trans = mc_utils.MCLabelTranslator() self.num_reduced_labels = self.label_trans.get_num_reduced_lbls() self.reduced_label_set = getattr(gen_cfg, 'reduced_label_set', False) self.use_label_smooth = getattr(gen_cfg, 'use_label_smooth', False) self.use_label_smooth_real = getattr(gen_cfg, 'use_label_smooth_real', self.use_label_smooth) self.use_label_smooth_pgt = getattr(gen_cfg, 'use_label_smooth_pgt', False) self.label_smooth_dia = getattr(gen_cfg, 'label_smooth_dia', 11) # Load MLP model. self.render_net = globals()[gen_cfg.mlp_model]( self.input_dim, viewdir_dim=self.input_dim_viewdir, style_dim=self.interm_style_dims, mask_dim=self.num_reduced_labels, out_channels_s=1, out_channels_c=self.final_feat_dim, **self.mlp_model_kwargs) # Camera sampler. self.camera_sampler_type = getattr(gen_cfg, 'camera_sampler_type', "random") assert self.camera_sampler_type in ['random', 'traditional'] self.camera_min_entropy = getattr(gen_cfg, 'camera_min_entropy', -1) self.camera_rej_avg_depth = getattr(gen_cfg, 'camera_rej_avg_depth', -1) self.cam_res = gen_cfg.cam_res self.crop_size = gen_cfg.crop_size print('Done with the GANcraft generator initialization.')
[docs] def custom_init(self): r"""Weight initialization of GANcraft components.""" self.blk_feats.data.uniform_(-1, 1) def init_func(m): if hasattr(m, 'weight'): nn.init.kaiming_normal_(m.weight.data, a=0.2, nonlinearity='leaky_relu') m.weight.data *= 0.5 if hasattr(m, 'bias') and m.bias is not None: m.bias.data.fill_(0.0) self.apply(init_func)
def _get_batch(self, batch_size, device): r"""Sample camera poses and perform ray-voxel intersection. Args: batch_size (int): Expected batch size of the current batch device (torch.device): Device on which the tensors should be stored """ with torch.no_grad(): voxel_id_batch = [] depth2_batch = [] raydirs_batch = [] cam_ori_t_batch = [] for b in range(batch_size): while True: # Rejection sampling. # Sample camera pose. if self.camera_sampler_type == 'random': cam_res = self.cam_res cam_ori_t, cam_dir_t, cam_up_t = camctl.rand_camera_pose_thridperson2(self.voxel) # ~24mm fov horizontal. cam_f = 0.5/np.tan(np.deg2rad(73/2) * (np.random.rand(1)*0.5+0.5)) * (cam_res[1]-1) cam_c = [(cam_res[0]-1)/2, (cam_res[1]-1)/2] cam_res_crop = [self.crop_size[0] + self.pad, self.crop_size[1] + self.pad] cam_c = mc_utils.rand_crop(cam_c, cam_res, cam_res_crop) elif self.camera_sampler_type == 'traditional': cam_res = self.cam_res cam_c = [(cam_res[0]-1)/2, (cam_res[1]-1)/2] dice = torch.rand(1).item() if dice > 0.5: cam_ori_t, cam_dir_t, cam_up_t, cam_f = \ camctl.rand_camera_pose_tour(self.voxel) cam_f = cam_f * (cam_res[1]-1) else: cam_ori_t, cam_dir_t, cam_up_t = \ camctl.rand_camera_pose_thridperson2(self.voxel) # ~24mm fov horizontal. cam_f = 0.5 / np.tan(np.deg2rad(73/2) * (np.random.rand(1)*0.5+0.5)) * (cam_res[1]-1) cam_res_crop = [self.crop_size[0] + self.pad, self.crop_size[1] + self.pad] cam_c = mc_utils.rand_crop(cam_c, cam_res, cam_res_crop) else: raise NotImplementedError( 'Unknown self.camera_sampler_type: {}'.format(self.camera_sampler_type)) # Run ray-voxel intersection test r"""Ray-voxel intersection CUDA kernel. Note: voxel_id = 0 and depth2 = NaN if there is no intersection along the ray Args: voxel_t (Y x 512 x 512 tensor, int32): Full 3D voxel of MC block IDs. cam_ori_t (3 tensor): Camera origin. cam_dir_t (3 tensor): Camera direction. cam_up_t (3 tensor): Camera up vector. cam_f (float): Camera focal length (in pixels). cam_c (list of 2 floats [x, y]): Camera optical center. img_dims (list of 2 ints [H, W]): Camera resolution. max_samples (int): Maximum number of blocks intersected along the ray before stopping. Returns: voxel_id ( img_dims[0] x img_dims[1] x max_samples x 1 tensor): IDs of intersected tensors along each ray depth2 (2 x img_dims[0] x img_dims[1] x max_samples x 1 tensor): Depths of entrance and exit points for each ray-voxel intersection. raydirs ( img_dims[0] x img_dims[1] x 1 x 3 tensor): The direction of each ray. """ voxel_id, depth2, raydirs = voxlib.ray_voxel_intersection_perspective( self.voxel.voxel_t, cam_ori_t, cam_dir_t, cam_up_t, cam_f, cam_c, cam_res_crop, self.num_blocks_early_stop) if self.camera_rej_avg_depth > 0: depth_map = depth2[0, :, :, 0, :] avg_depth = torch.mean(depth_map[~torch.isnan(depth_map)]) if avg_depth < self.camera_rej_avg_depth: continue # Reject low entropy. if self.camera_min_entropy > 0: # Check entropy. maskcnt = torch.bincount( torch.flatten(voxel_id[:, :, 0, 0]), weights=None, minlength=680).float() / \ (voxel_id.size(0)*voxel_id.size(1)) maskentropy = -torch.sum(maskcnt * torch.log(maskcnt+1e-10)) if maskentropy < self.camera_min_entropy: continue break voxel_id_batch.append(voxel_id) depth2_batch.append(depth2) raydirs_batch.append(raydirs) cam_ori_t_batch.append(cam_ori_t) voxel_id = torch.stack(voxel_id_batch, dim=0) depth2 = torch.stack(depth2_batch, dim=0) raydirs = torch.stack(raydirs_batch, dim=0) cam_ori_t = torch.stack(cam_ori_t_batch, dim=0).to(device) cam_poses = None return voxel_id, depth2, raydirs, cam_ori_t, cam_poses
[docs] def get_pseudo_gt(self, pseudo_gen, voxel_id, z=None, style_img=None, resize_512=True, deterministic=False): r"""Evaluating img2img network to obtain pseudo-ground truth images. Args: pseudo_gen (callable): Function converting mask to image using img2img network. voxel_id (N x img_dims[0] x img_dims[1] x max_samples x 1 tensor): IDs of intersected tensors along each ray. z (N x C tensor): Optional style code passed to pseudo_gen. style_img (N x 3 x H x W tensor): Optional style image passed to pseudo_gen. resize_512 (bool): If True, evaluate pseudo_gen at 512x512 regardless of input resolution. deterministic (bool): If True, disable stochastic label mapping. """ with torch.no_grad(): mc_mask = voxel_id[:, :, :, 0, :].permute(0, 3, 1, 2).long() coco_mask = self.label_trans.mc2coco(mc_mask) - 1 coco_mask[coco_mask < 0] = 183 if not deterministic: # Stochastic mapping dice = torch.rand(1).item() if dice > 0.5 and dice < 0.9: coco_mask[coco_mask == self.label_trans.gglbl2ggid('sky')] = self.label_trans.gglbl2ggid('clouds') elif dice >= 0.9: coco_mask[coco_mask == self.label_trans.gglbl2ggid('sky')] = self.label_trans.gglbl2ggid('fog') dice = torch.rand(1).item() if dice > 0.33 and dice < 0.66: coco_mask[coco_mask == self.label_trans.gglbl2ggid('water')] = self.label_trans.gglbl2ggid('sea') elif dice >= 0.66: coco_mask[coco_mask == self.label_trans.gglbl2ggid('water')] = self.label_trans.gglbl2ggid('river') fake_masks = torch.zeros([coco_mask.size(0), 185, coco_mask.size(2), coco_mask.size(3)], dtype=torch.half, device=voxel_id.device) fake_masks.scatter_(1, coco_mask, 1.0) if self.use_label_smooth_pgt: fake_masks = mc_utils.segmask_smooth(fake_masks, kernel_size=self.label_smooth_dia) if self.pad > 0: fake_masks = fake_masks[:, :, self.pad//2:-self.pad//2, self.pad//2:-self.pad//2] # Generate pseudo GT using GauGAN. if resize_512: fake_masks_512 = F.interpolate(fake_masks, size=[512, 512], mode='nearest') else: fake_masks_512 = fake_masks pseudo_real_img = pseudo_gen(fake_masks_512, z=z, style_img=style_img) # NaN Inf Guard. NaN can occure on Volta GPUs. nan_mask = torch.isnan(pseudo_real_img) inf_mask = torch.isinf(pseudo_real_img) pseudo_real_img[nan_mask | inf_mask] = 0.0 if resize_512: pseudo_real_img = F.interpolate( pseudo_real_img, size=[fake_masks.size(2), fake_masks.size(3)], mode='area') pseudo_real_img = torch.clamp(pseudo_real_img, -1, 1) return pseudo_real_img, fake_masks
[docs] def sample_camera(self, data, pseudo_gen): r"""Sample camera randomly and precompute everything used by both Gen and Dis. Args: data (dict): images (N x 3 x H x W tensor) : Real images label (N x C2 x H x W tensor) : Segmentation map pseudo_gen (callable): Function converting mask to image using img2img network. Returns: ret (dict): voxel_id (N x H x W x max_samples x 1 tensor): IDs of intersected tensors along each ray. depth2 (N x 2 x H x W x max_samples x 1 tensor): Depths of entrance and exit points for each ray-voxel intersection. raydirs (N x H x W x 1 x 3 tensor): The direction of each ray. cam_ori_t (N x 3 tensor): Camera origins. pseudo_real_img (N x 3 x H x W tensor): Pseudo-ground truth image. real_masks (N x C3 x H x W tensor): One-hot segmentation map for real images, with translated labels. fake_masks (N x C3 x H x W tensor): One-hot segmentation map for sampled camera views. """ device = torch.device('cuda') batch_size = data['images'].size(0) # ================ Assemble a batch ================== # Requires: voxel_id, depth2, raydirs, cam_ori_t. voxel_id, depth2, raydirs, cam_ori_t, _ = self._get_batch(batch_size, device) ret = {'voxel_id': voxel_id, 'depth2': depth2, 'raydirs': raydirs, 'cam_ori_t': cam_ori_t} if pseudo_gen is not None: pseudo_real_img, _ = self.get_pseudo_gt(pseudo_gen, voxel_id) ret['pseudo_real_img'] = pseudo_real_img.float() # =============== Mask translation ================ real_masks = data['label'] if self.reduced_label_set: # Translate fake mask (directly from mcid). # convert unrecognized labels to 'dirt'. # N C H W [1 1 80 80] reduce_fake_mask = self.label_trans.mc2reduced( voxel_id[:, :, :, 0, :].permute(0, 3, 1, 2).long(), ign2dirt=True) reduce_fake_mask_onehot = torch.zeros([ reduce_fake_mask.size(0), self.num_reduced_labels, reduce_fake_mask.size(2), reduce_fake_mask.size(3)], dtype=torch.float, device=device) reduce_fake_mask_onehot.scatter_(1, reduce_fake_mask, 1.0) fake_masks = reduce_fake_mask_onehot if self.pad != 0: fake_masks = fake_masks[:, :, self.pad//2:-self.pad//2, self.pad//2:-self.pad//2] # Translate real mask (data['label']), which is onehot. real_masks_idx = torch.argmax(real_masks, dim=1, keepdim=True) real_masks_idx[real_masks_idx > 182] = 182 reduced_real_mask = self.label_trans.coco2reduced(real_masks_idx) reduced_real_mask_onehot = torch.zeros([ reduced_real_mask.size(0), self.num_reduced_labels, reduced_real_mask.size(2), reduced_real_mask.size(3)], dtype=torch.float, device=device) reduced_real_mask_onehot.scatter_(1, reduced_real_mask, 1.0) real_masks = reduced_real_mask_onehot # Mask smoothing. if self.use_label_smooth: fake_masks = mc_utils.segmask_smooth(fake_masks, kernel_size=self.label_smooth_dia) if self.use_label_smooth_real: real_masks = mc_utils.segmask_smooth(real_masks, kernel_size=self.label_smooth_dia) ret['real_masks'] = real_masks ret['fake_masks'] = fake_masks return ret
[docs] def forward(self, data, random_style=False): r"""GANcraft Generator forward. Args: data (dict): images (N x 3 x H x W tensor) : Real images voxel_id (N x H x W x max_samples x 1 tensor): IDs of intersected tensors along each ray. depth2 (N x 2 x H x W x max_samples x 1 tensor): Depths of entrance and exit points for each ray-voxel intersection. raydirs (N x H x W x 1 x 3 tensor): The direction of each ray. cam_ori_t (N x 3 tensor): Camera origins. random_style (bool): Whether to sample a random style vector. Returns: output (dict): fake_images (N x 3 x H x W tensor): fake images mu (N x C1 tensor): mean vectors logvar (N x C1 tensor): log-variance vectors """ device = torch.device('cuda') batch_size = data['images'].size(0) # ================ Assemble a batch ================== # Requires: voxel_id, depth2, raydirs, cam_ori_t. voxel_id, depth2, raydirs, cam_ori_t = data['voxel_id'], data['depth2'], data['raydirs'], data['cam_ori_t'] if 'pseudo_real_img' in data: pseudo_real_img = data['pseudo_real_img'] z, mu, logvar = None, None, None if random_style: if self.style_dims > 0: z = torch.randn(batch_size, self.style_dims, dtype=torch.float32, device=device) else: if self.style_encoder is None: # ================ Get Style Code ================= if self.style_dims > 0: z = torch.randn(batch_size, self.style_dims, dtype=torch.float32, device=device) else: mu, logvar, z = self.style_encoder(pseudo_real_img) # ================ Network Forward ================ # Forward StyleNet if self.style_net is not None: z = self.style_net(z) # Forward per-pixel net. net_out, new_dists, weights, total_weights_raw, rand_depth, net_out_s, net_out_c, skynet_out_c, nosky_mask, \ sky_mask, sky_only_mask, new_idx = self._forward_perpix( self.blk_feats, voxel_id, depth2, raydirs, cam_ori_t, z) # Forward global net. fake_images, fake_images_raw = self._forward_global(net_out, z) if self.pad != 0: fake_images = fake_images[:, :, self.pad//2:-self.pad//2, self.pad//2:-self.pad//2] # =============== Arrange Return Values ================ output = {} output['fake_images'] = fake_images output['mu'] = mu output['logvar'] = logvar return output
[docs] def inference(self, output_dir, camera_mode, style_img_path=None, seed=1, pad=30, num_samples=40, num_blocks_early_stop=6, sample_depth=3, tile_size=128, resolution_hw=[540, 960], cam_ang=72, cam_maxstep=10): r"""Compute result images according to the provided camera trajectory and save the results in the specified folder. The full image is evaluated in multiple tiles to save memory. Args: output_dir (str): Where should the results be stored. camera_mode (int): Which camera trajectory to use. style_img_path (str): Path to the style-conditioning image. seed (int): Random seed (controls style when style_image_path is not specified). pad (int): Pixels to remove from the image tiles before stitching. Should be equal or larger than the receptive field of the CNN to avoid border artifact. num_samples (int): Number of samples per ray (different from training). num_blocks_early_stop (int): Max number of intersected boxes per ray before stopping (different from training). sample_depth (float): Max distance traveled through boxes before stopping (different from training). tile_size (int): Max size of a tile in pixels. resolution_hw (list [H, W]): Resolution of the output image. cam_ang (float): Horizontal FOV of the camera (may be adjusted by the camera controller). cam_maxstep (int): Number of frames sampled from the camera trajectory. """ def write_img(path, img, rgb_input=False): img = ((img*0.5+0.5)*255).detach().cpu().numpy().astype(np.uint8) img = img[0].transpose(1, 2, 0) if rgb_input: img = img[..., [2, 1, 0]] cv2.imwrite(path, img, [cv2.IMWRITE_PNG_COMPRESSION, 4]) return img[..., ::-1] def read_img(path): img = cv2.imread(path).astype(np.float32)[..., [2, 1, 0]].transpose(2, 0, 1) / 255 img = img * 2 - 1 img = torch.from_numpy(img) print('Saving to', output_dir) # Use provided random seed. device = torch.device('cuda') rng_cuda = torch.Generator(device=device) rng_cuda = rng_cuda.manual_seed(seed) torch.manual_seed(seed) torch.cuda.manual_seed(seed) self.pad = pad self.num_samples = num_samples self.num_blocks_early_stop = num_blocks_early_stop self.sample_depth = sample_depth self.coarse_deterministic_sampling = True self.crop_size = resolution_hw self.cam_res = [self.crop_size[0]+self.pad, self.crop_size[1]+self.pad] self.use_label_smooth_pgt = False # Make output dirs. gancraft_outputs_dir = os.path.join(output_dir, 'gancraft_outputs') os.makedirs(gancraft_outputs_dir, exist_ok=True) vis_masks_dir = os.path.join(output_dir, 'vis_masks') os.makedirs(vis_masks_dir, exist_ok=True) fout = imageio.get_writer(gancraft_outputs_dir + '.mp4', fps=1) fout_cat = imageio.get_writer(gancraft_outputs_dir + '-vis_masks.mp4', fps=1) evalcamctl = camctl.EvalCameraController( self.voxel, maxstep=cam_maxstep, pattern=camera_mode, cam_ang=cam_ang, smooth_decay_multiplier=150/cam_maxstep) # Get output style. if style_img_path is None: z = torch.empty(1, self.style_dims, dtype=torch.float32, device=device) z.normal_(generator=rng_cuda) else: style_img = read_img(style_img_path) style_img = style_img.to(device).unsqueeze(0) mu, logvar, z = self.style_encoder(style_img) z = self.style_net(z) # Generate required output images. for id, (cam_ori_t, cam_dir_t, cam_up_t, cam_f) in enumerate(evalcamctl): print('Rendering frame', id) cam_f = cam_f * (self.crop_size[1]-1) # So that the view is not depending on the padding cam_c = [(self.cam_res[0]-1)/2, (self.cam_res[1]-1)/2] voxel_id, depth2, raydirs = voxlib.ray_voxel_intersection_perspective( self.voxel.voxel_t, cam_ori_t, cam_dir_t, cam_up_t, cam_f, cam_c, self.cam_res, self.num_blocks_early_stop) voxel_id = voxel_id.unsqueeze(0) depth2 = depth2.unsqueeze(0) raydirs = raydirs.unsqueeze(0) cam_ori_t = cam_ori_t.unsqueeze(0).to(device) # Save 3D voxel rendering. mc_rgb = self.label_trans.mc_color(voxel_id[0, :, :, 0, 0].cpu().numpy()) # Diffused shading, co-located light. first_intersection_depth = depth2[:, 0, :, :, 0, None, :] # [1, 542, 542, 1, 1]. first_intersection_point = raydirs * first_intersection_depth + cam_ori_t[:, None, None, None, :] fip_local_coords = torch.remainder(first_intersection_point, 1.0) fip_wall_proximity = torch.minimum(fip_local_coords, 1.0-fip_local_coords) fip_wall_orientation = torch.argmin(fip_wall_proximity, dim=-1, keepdim=False) # 0: [1,0,0]; 1: [0,1,0]; 2: [0,0,1] lut = torch.tensor([[1, 0, 0], [0, 1, 0], [0, 0, 1]], dtype=torch.float32, device=fip_wall_orientation.device) fip_normal = lut[fip_wall_orientation] # [1, 542, 542, 1, 3] diffuse_shade = torch.abs(torch.sum(fip_normal * raydirs, dim=-1)) mc_rgb = (mc_rgb.astype(np.float) / 255) ** 2.2 mc_rgb = mc_rgb * diffuse_shade[0, :, :, :].cpu().numpy() mc_rgb = (mc_rgb ** (1/2.2)) * 255 mc_rgb = mc_rgb.astype(np.uint8) if self.pad > 0: mc_rgb = mc_rgb[self.pad//2:-self.pad//2, self.pad//2:-self.pad//2] cv2.imwrite(os.path.join(vis_masks_dir, '{:05d}.png'.format(id)), mc_rgb, [cv2.IMWRITE_PNG_COMPRESSION, 4]) # Tiled eval of GANcraft. voxel_id_all = voxel_id depth2_all = depth2 raydirs_all = raydirs # Evaluate sky in advance to get a consistent sky in the semi-transparent region. if self.sky_global_avgpool: sky_raydirs_in = raydirs.expand(-1, -1, -1, 1, -1).contiguous() sky_raydirs_in = voxlib.positional_encoding( sky_raydirs_in, self.pe_params_sky[0], -1, self.pe_params_sky[1]) skynet_out_c = self.sky_net(sky_raydirs_in, z) sky_avg = torch.mean(skynet_out_c, dim=[1, 2], keepdim=True) self.sky_avg = sky_avg num_strips_h = (self.cam_res[0]-self.pad+tile_size-1)//tile_size num_strips_w = (self.cam_res[1]-self.pad+tile_size-1)//tile_size fake_images_chunks_v = [] # For each horizontal strip. for strip_id_h in range(num_strips_h): strip_begin_h = strip_id_h * tile_size strip_end_h = np.minimum(strip_id_h * tile_size + tile_size + self.pad, self.cam_res[0]) # For each vertical strip. fake_images_chunks_h = [] for strip_id_w in range(num_strips_w): strip_begin_w = strip_id_w * tile_size strip_end_w = np.minimum(strip_id_w * tile_size + tile_size + self.pad, self.cam_res[1]) voxel_id = voxel_id_all[:, strip_begin_h:strip_end_h, strip_begin_w:strip_end_w, :, :] depth2 = depth2_all[:, :, strip_begin_h:strip_end_h, strip_begin_w:strip_end_w, :, :] raydirs = raydirs_all[:, strip_begin_h:strip_end_h, strip_begin_w:strip_end_w, :, :] net_out, new_dists, weights, total_weights_raw, rand_depth, net_out_s, net_out_c, skynet_out_c, \ nosky_mask, sky_mask, sky_only_mask, new_idx = self._forward_perpix( self.blk_feats, voxel_id, depth2, raydirs, cam_ori_t, z) fake_images, _ = self._forward_global(net_out, z) if self.pad != 0: fake_images = fake_images[:, :, self.pad//2:-self.pad//2, self.pad//2:-self.pad//2] fake_images_chunks_h.append(fake_images) fake_images_h = torch.cat(fake_images_chunks_h, dim=-1) fake_images_chunks_v.append(fake_images_h) fake_images = torch.cat(fake_images_chunks_v, dim=-2) rgb = write_img(os.path.join(gancraft_outputs_dir, '{:05d}.png'.format(id)), fake_images, rgb_input=True) fout.append_data(rgb) fout_cat.append_data(np.concatenate((mc_rgb[..., ::-1], rgb), axis=1)) fout.close() fout_cat.close()