Skip to content

Commit

Permalink
Merge pull request #28 from lijianshe02/master
Browse files Browse the repository at this point in the history
add remove duplicate frames in DAIN
  • Loading branch information
qingqing01 authored Sep 28, 2020
2 parents 5815414 + c02f407 commit 4cd9a0b
Show file tree
Hide file tree
Showing 3 changed files with 91 additions and 58 deletions.
5 changes: 4 additions & 1 deletion applications/DAIN/my_args.py
Original file line number Diff line number Diff line change
Expand Up @@ -90,4 +90,7 @@
type=bool,
help='use cuda or not')
parser.add_argument('--use_cudnn', default=1, type=int, help='use cudnn or not')

parser.add_argument('--remove_duplicates',
default=True,
type=bool,
help='remove duplicate frames or not')
112 changes: 55 additions & 57 deletions applications/DAIN/predict.py
Original file line number Diff line number Diff line change
Expand Up @@ -80,7 +80,8 @@ def __init__(self,
video_path,
use_gpu=True,
key_frame_thread=0.,
output_path='output'):
output_path='output',
remove_duplicates=True):
self.video_path = video_path
self.output_path = os.path.join(output_path, 'DAIN')
if model_path is None:
Expand Down Expand Up @@ -138,6 +139,8 @@ def run(self):
end = time.time()

frames = sorted(glob.glob(os.path.join(out_path, '*.png')))
if remove_duplicates:
frames = remove_duplicates(out_path)

img = imread(frames[0])

Expand Down Expand Up @@ -199,58 +202,51 @@ def run(self):
X0 = img_first.astype('float32').transpose((2, 0, 1)) / 255
X1 = img_second.astype('float32').transpose((2, 0, 1)) / 255

if key_frame:
y_ = [
np.transpose(255.0 * X0.clip(0, 1.0), (1, 2, 0))
for i in range(num_frames)
]
else:
assert (X0.shape[1] == X1.shape[1])
assert (X0.shape[2] == X1.shape[2])

X0 = np.pad(X0, ((0,0), (padding_top, padding_bottom), \
(padding_left, padding_right)), mode='edge')
X1 = np.pad(X1, ((0,0), (padding_top, padding_bottom), \
(padding_left, padding_right)), mode='edge')

X0 = np.expand_dims(X0, axis=0)
X1 = np.expand_dims(X1, axis=0)

X0 = np.expand_dims(X0, axis=0)
X1 = np.expand_dims(X1, axis=0)

X = np.concatenate((X0, X1), axis=0)

proc_end = time.time()
o = self.exe.run(self.program,
fetch_list=self.fetch_targets,
feed={"image": X})

y_ = o[0]

proc_timer.update(time.time() - proc_end)
tot_timer.update(time.time() - end)
end = time.time()

y_ = [
np.transpose(
255.0 * item.clip(
0, 1.0)[0, :,
padding_top:padding_top + int_height,
padding_left:padding_left + int_width],
(1, 2, 0)) for item in y_
]
time_offsets = [
kk * timestep for kk in range(1, 1 + num_frames, 1)
]

count = 1
for item, time_offset in zip(y_, time_offsets):
out_dir = os.path.join(
frame_path_interpolated, vidname,
"{:0>6d}_{:0>4d}.png".format(i, count))
count = count + 1
imsave(out_dir, np.round(item).astype(np.uint8))
assert (X0.shape[1] == X1.shape[1])
assert (X0.shape[2] == X1.shape[2])

X0 = np.pad(X0, ((0,0), (padding_top, padding_bottom), \
(padding_left, padding_right)), mode='edge')
X1 = np.pad(X1, ((0,0), (padding_top, padding_bottom), \
(padding_left, padding_right)), mode='edge')

X0 = np.expand_dims(X0, axis=0)
X1 = np.expand_dims(X1, axis=0)

X0 = np.expand_dims(X0, axis=0)
X1 = np.expand_dims(X1, axis=0)

X = np.concatenate((X0, X1), axis=0)

proc_end = time.time()
o = self.exe.run(self.program,
fetch_list=self.fetch_targets,
feed={"image": X})

y_ = o[0]

proc_timer.update(time.time() - proc_end)
tot_timer.update(time.time() - end)
end = time.time()

y_ = [
np.transpose(
255.0 * item.clip(
0, 1.0)[0, :, padding_top:padding_top + int_height,
padding_left:padding_left + int_width],
(1, 2, 0)) for item in y_
]
time_offsets = [
kk * timestep for kk in range(1, 1 + num_frames, 1)
]

count = 1
for item, time_offset in zip(y_, time_offsets):
out_dir = os.path.join(
frame_path_interpolated, vidname,
"{:0>6d}_{:0>4d}.png".format(i, count))
count = count + 1
imsave(out_dir, np.round(item).astype(np.uint8))

num_frames = int(1.0 / timestep) - 1

Expand All @@ -266,14 +262,16 @@ def run(self):
vidname + '.mp4')
if os.path.exists(video_pattern_output):
os.remove(video_pattern_output)
frames2video(frame_pattern_combined, video_pattern_output,
r2)
frames2video(frame_pattern_combined, video_pattern_output, r2)

return frame_pattern_combined, video_pattern_output


if __name__ == '__main__':
args = parser.parse_args()
predictor = VideoFrameInterp(args.time_step, args.saved_model,
args.video_path, args.output_path)
predictor = VideoFrameInterp(args.time_step,
args.saved_model,
args.video_path,
args.output_path,
remove_duplicates=args.remove_duplicates)
predictor.run()
32 changes: 32 additions & 0 deletions applications/DAIN/util.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
import os, sys
import glob
import shutil
import cv2


class AverageMeter(object):
Expand Down Expand Up @@ -44,3 +45,34 @@ def combine_frames(input, interpolated, combined, num_frames):
except Exception as e:
print(e)
print(len(frames2), num_frames, i, k, i * num_frames + k)


def remove_duplicates(paths):
def dhash(image, hash_size=8):
gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
resized = cv2.resize(gray, (hash_size + 1, hash_size))
diff = resized[:, 1:] > resized[:, :-1]
return sum([2**i for (i, v) in enumerate(diff.flatten()) if v])

hashes = {}
image_paths = sorted(glob.glob(os.path.join(paths, '*.png')))
for image_path in image_paths:
image = cv2.imread(image_path)
h = dhash(image)
p = hashes.get(h, [])
p.append(image_path)
hashes[h] = p

for (h, hashed_paths) in hashes.items():
if len(hashed_paths) > 1:
for p in hashed_paths[1:]:
os.remove(p)

frames = sorted(glob.glob(os.path.join(paths, '*.png')))
for fid, frame in enumerate(frames):
new_name = '{:08d}'.format(fid) + '.png'
new_name = os.path.join(paths, new_name)
os.rename(frame, new_name)

frames = sorted(glob.glob(os.path.join(paths, '*.png')))
return frames

0 comments on commit 4cd9a0b

Please sign in to comment.