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 )