def run(self): outputdir = self.output outputdir_in = os.path.join(outputdir, 'input/') os.makedirs(outputdir_in, exist_ok=True) outputdir_out = os.path.join(outputdir, 'output/') os.makedirs(outputdir_out, exist_ok=True) # Prepare reference images if self.colorization: if self.reference_dir is not None: import glob ext_list = ['png', 'jpg', 'bmp'] reference_files = [] for ext in ext_list: reference_files += glob.glob(self.reference_dir + '/*.' + ext, recursive=True) aspect_mean = 0 minedge_dim = 256 refs = [] for v in reference_files: refimg = Image.open(v).convert('RGB') w, h = refimg.size aspect_mean += w / h refs.append(refimg) aspect_mean /= len(reference_files) target_w = int(256 * aspect_mean) if aspect_mean > 1 else 256 target_h = 256 if aspect_mean >= 1 else int(256 / aspect_mean) refimgs = [] for i, v in enumerate(refs): refimg = utils.addMergin(v, target_w=target_w, target_h=target_h) refimg = np.array(refimg).astype('float32').transpose( 2, 0, 1) / 255.0 refimgs.append(refimg) refimgs = paddle.to_tensor(np.array(refimgs).astype('float32')) refimgs = paddle.unsqueeze(refimgs, 0) # Load video cap = cv2.VideoCapture(self.input) nframes = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) v_w = cap.get(cv2.CAP_PROP_FRAME_WIDTH) v_h = cap.get(cv2.CAP_PROP_FRAME_HEIGHT) minwh = min(v_w, v_h) scale = 1 if minwh != self.mindim: scale = self.mindim / minwh t_w = round(v_w * scale / 16.) * 16 t_h = round(v_h * scale / 16.) * 16 fps = cap.get(cv2.CAP_PROP_FPS) pbar = tqdm(total=nframes) block = 5 # Process with paddle.no_grad(): it = 0 while True: frame_pos = it * block if frame_pos >= nframes: break cap.set(cv2.CAP_PROP_POS_FRAMES, frame_pos) if block >= nframes - frame_pos: proc_g = nframes - frame_pos else: proc_g = block input = None gtC = None for i in range(proc_g): index = frame_pos + i _, frame = cap.read() frame = cv2.resize(frame, (t_w, t_h)) nchannels = frame.shape[2] if nchannels == 1 or self.colorization: frame_l = cv2.cvtColor(frame, cv2.COLOR_RGB2GRAY) cv2.imwrite(outputdir_in + '%07d.png' % index, frame_l) frame_l = paddle.to_tensor(frame_l.astype('float32')) frame_l = paddle.reshape( frame_l, [frame_l.shape[0], frame_l.shape[1], 1]) frame_l = paddle.transpose(frame_l, [2, 0, 1]) frame_l /= 255. frame_l = paddle.reshape(frame_l, [ 1, frame_l.shape[0], 1, frame_l.shape[1], frame_l.shape[2] ]) elif nchannels == 3: cv2.imwrite(outputdir_in + '%07d.png' % index, frame) frame = frame[:, :, ::-1] ## BGR -> RGB frame_l, frame_ab = utils.convertRGB2LABTensor(frame) frame_l = frame_l.transpose([2, 0, 1]) frame_ab = frame_ab.transpose([2, 0, 1]) frame_l = frame_l.reshape([ 1, frame_l.shape[0], 1, frame_l.shape[1], frame_l.shape[2] ]) frame_ab = frame_ab.reshape([ 1, frame_ab.shape[0], 1, frame_ab.shape[1], frame_ab.shape[2] ]) if input is not None: paddle.concat((input, frame_l), 2) input = frame_l if i == 0 else paddle.concat( (input, frame_l), 2) if nchannels == 3 and not self.colorization: gtC = frame_ab if i == 0 else paddle.concat( (gtC, frame_ab), 2) input = paddle.to_tensor(input) output_l = self.modelR(input) # [B, C, T, H, W] # Save restoration output without colorization when using the option [--disable_colorization] if not self.colorization: for i in range(proc_g): index = frame_pos + i if nchannels == 3: out_l = output_l.detach()[0, :, i] out_ab = gtC[0, :, i] out = paddle.concat( (out_l, out_ab), axis=0).detach().numpy().transpose((1, 2, 0)) out = Image.fromarray( np.uint8(utils.convertLAB2RGB(out) * 255)) out.save(outputdir_out + '%07d.png' % (index)) else: raise ValueError('channels of imag3 must be 3!') # Perform colorization else: if self.reference_dir is None: output_ab = self.modelC(output_l) else: output_ab = self.modelC(output_l, refimgs) output_l = output_l.detach() output_ab = output_ab.detach() for i in range(proc_g): index = frame_pos + i out_l = output_l[0, :, i, :, :] out_c = output_ab[0, :, i, :, :] output = paddle.concat((out_l, out_c), axis=0).numpy().transpose( (1, 2, 0)) output = Image.fromarray( np.uint8(utils.convertLAB2RGB(output) * 255)) output.save(outputdir_out + '%07d.png' % index) it = it + 1 pbar.update(proc_g) # Save result videos outfile = os.path.join(outputdir, self.input.split('/')[-1].split('.')[0]) cmd = 'ffmpeg -y -r %d -i %s%%07d.png -vcodec libx264 -pix_fmt yuv420p -r %d %s_in.mp4' % ( fps, outputdir_in, fps, outfile) subprocess.call(cmd, shell=True) cmd = 'ffmpeg -y -r %d -i %s%%07d.png -vcodec libx264 -pix_fmt yuv420p -r %d %s_out.mp4' % ( fps, outputdir_out, fps, outfile) subprocess.call(cmd, shell=True) cmd = 'ffmpeg -y -i %s_in.mp4 -vf "[in] pad=2.01*iw:ih [left];movie=%s_out.mp4[right];[left][right] overlay=main_w/2:0,scale=2*iw/2:2*ih/2[out]" %s_comp.mp4' % ( outfile, outfile, outfile) subprocess.call(cmd, shell=True) cap.release() pbar.close() return outputdir_out, '%s_out.mp4' % outfile
gtC = frame_ab if i==0 else torch.cat( (gtC, frame_ab), 2 ) input = input.to( device ) # Perform restoration output_l = modelR( input ) # [B, C, T, H, W] # Save restoration output without colorization when using the option [--disable_colorization] if opt.disable_colorization: for i in range( proc_g ): index = frame_pos + i if nchannels==3: out_l = output_l.detach()[0,:,i].cpu() out_ab = gtC[0,:,i].cpu() out = torch.cat((out_l, out_ab),dim=0).detach().numpy().transpose((1, 2, 0)) out = Image.fromarray( np.uint8( utils.convertLAB2RGB( out )*255 ) ) out.save( outputdir_out+'%07d.png'%(index) ) else: save_image( output_l.detach()[0,:,i], outputdir_out+'%07d.png'%(index), nrow=1 ) # Perform colorization else: if opt.reference_dir=='none': output_ab = modelC( output_l ) else: output_ab = modelC( output_l, refimgs ) output_l = output_l.detach().cpu() output_ab = output_ab.detach().cpu() # Save output frames of restoration with colorization for i in range( proc_g ): index = frame_pos + i
# (torch.Size([1, 1, 10, 320, 432]), torch.Size([1, 2, 10, 320, 432])) # (Pdb) output_ab.min(), output_ab.max() # (tensor(0.3863), tensor(0.6358)) # refimgs.size() # torch.Size([1, 4, 3, 256, 341]) # Save output frames of restoration with colorization for i in range(proc_g): index = frame_pos + i out_l = output_l[0, :, i, :, :] out_c = output_ab[0, :, i, :, :] output = torch.cat((out_l, out_c), dim=0).numpy().transpose( (1, 2, 0)) output = Image.fromarray( np.uint8(utils.convertLAB2RGB(output) * 255)) output.save(outputdir_out + '%07d.png' % index) # pdb.set_trace() it = it + 1 pbar.update(proc_g) # # Save result videos # outfile = opt.input.split('/')[-1].split('.')[0] # cmd = 'ffmpeg -y -r %d -i %s%%07d.png -vcodec libx264 -pix_fmt yuv420p -r %d %s_in.mp4' % (fps, outputdir_in, fps, outfile ) # subprocess.call( cmd, shell=True ) # cmd = 'ffmpeg -y -r %d -i %s%%07d.png -vcodec libx264 -pix_fmt yuv420p -r %d %s_out.mp4' % (fps, outputdir_out, fps, outfile ) # subprocess.call( cmd, shell=True ) # cmd = 'ffmpeg -y -i %s_in.mp4 -vf "[in] pad=2.01*iw:ih [left];movie=%s_out.mp4[right];[left][right] overlay=main_w/2:0,scale=2*iw/2:2*ih/2[out]" %s_comp.mp4' % ( outfile, outfile, outfile ) # subprocess.call( cmd, shell=True )