Mirrors/ECCV2022-RIFE
1
0
Fork 0
mirror of https://github.com/hzwer/ECCV2022-RIFE.git synced 2026-02-24 04:19:41 +01:00
Code Issues Packages Projects Releases Wiki Activity

Remove read buffer

Browse Source
This commit is contained in:
hzwer
2020-12-05 00:29:47 +08:00
parent 8286709701
commit 7ea17e61e8
1 changed files with 17 additions and 27 deletions
Download Patch File Download Diff File Expand all files Collapse all files

44
inference_video.py
View File

@@ -101,12 +101,12 @@ if args.png:
if not os.path.exists('vid_out'):
os.mkdir('vid_out')
else:
vid_out = cv2.VideoWriter('{}_{}X_{}fps.{}'.format(video_path_wo_ext, (2**args.exp), int(np.round(args.fps)), args.ext), fourcc, args.fps, (w, h))
vid_out = cv2.VideoWriter('{}_{}X_{}fps.{}'.format(video_path_wo_ext, (2 ** args.exp), int(np.round(args.fps)), args.ext), fourcc, args.fps, (w, h))
def clear_write_buffer(user_args, write_buffer):
def clear_buffer(user_args, buffer):
cnt = 0
while True:
item = write_buffer.get()
item = buffer.get()
if item is None:
break
if user_args.png:
@@ -115,15 +115,6 @@ def clear_write_buffer(user_args, write_buffer):
else:
vid_out.write(item[:, :, ::-1])
def build_read_buffer(user_args, read_buffer, videogen):
for frame in videogen:
if not user_args.img is None:
frame = cv2.imread(os.path.join(user_args.img, frame))[:, :, ::-1].copy()
if user_args.montage:
frame = frame[:, left: left + w]
read_buffer.put(frame)
read_buffer.put(None)
def make_inference(I0, I1, exp):
global model
middle = model.inference(I0, I1)
@@ -143,15 +134,14 @@ pbar = tqdm(total=tot_frame)
skip_frame = 1
if args.montage:
lastframe = lastframe[:, left: left + w]
write_buffer = Queue(maxsize=100)
read_buffer = Queue(maxsize=100)
_thread.start_new_thread(build_read_buffer, (args, read_buffer, videogen))
_thread.start_new_thread(clear_write_buffer, (args, write_buffer))
buffer = Queue()
_thread.start_new_thread(clear_buffer, (args, buffer))
while True:
frame = read_buffer.get()
if frame is None:
break
for frame in videogen:
if not args.img is None:
frame = cv2.imread(os.path.join(args.img, frame))[:, :, ::-1].copy()
if args.montage:
frame = frame[:, left: left + w]
I0 = torch.from_numpy(np.transpose(lastframe, (2,0,1))).to(device, non_blocking=True).unsqueeze(0).float() / 255.
I1 = torch.from_numpy(np.transpose(frame, (2,0,1))).to(device, non_blocking=True).unsqueeze(0).float() / 255.
I0 = F.pad(I0, padding)
@@ -171,23 +161,23 @@ while True:
else:
output = make_inference(I0, I1, args.exp)
if args.montage:
write_buffer.put(np.concatenate((lastframe, lastframe), 1))
buffer.put(np.concatenate((lastframe, lastframe), 1))
for mid in output:
mid = (((mid[0] * 255.).byte().cpu().detach().numpy().transpose(1, 2, 0)))
write_buffer.put(np.concatenate((lastframe, mid[:h, :w]), 1))
buffer.put(np.concatenate((lastframe, mid[:h, :w]), 1))
else:
write_buffer.put(lastframe)
buffer.put(lastframe)
for mid in output:
mid = (((mid[0] * 255.).byte().cpu().detach().numpy().transpose(1, 2, 0)))
write_buffer.put(mid[:h, :w])
buffer.put(mid[:h, :w])
pbar.update(1)
lastframe = frame
if args.montage:
write_buffer.put(np.concatenate((lastframe, lastframe), 1))
buffer.put(np.concatenate((lastframe, lastframe), 1))
else:
write_buffer.put(lastframe)
buffer.put(lastframe)
import time
while(not write_buffer.empty()):
while(not buffer.empty()):
time.sleep(0.1)
pbar.close()
if not vid_out is None: