|
import os |
|
import cv2 |
|
import numpy as np |
|
import torch |
|
import bisect |
|
import shutil |
|
|
|
if torch.backends.mps.is_available(): |
|
device = "mps" |
|
|
|
elif torch.cuda.is_available(): |
|
device = "cuda" |
|
else: |
|
device = "cpu" |
|
|
|
def init_frame_interpolation_model(): |
|
print("Initializing frame interpolation model") |
|
checkpoint_name = os.path.join("./pretrained_model/film_net_fp16.pt") |
|
|
|
model = torch.load(checkpoint_name, map_location='cpu') |
|
model.eval() |
|
model = model.half() |
|
model = model.to(device=device) |
|
return model |
|
|
|
|
|
def batch_images_interpolation_tool(input_file, model, fps, inter_frames=1): |
|
|
|
image_save_dir = input_file + '_tmp' |
|
os.makedirs(image_save_dir, exist_ok=True) |
|
|
|
input_img_list = os.listdir(input_file) |
|
input_img_list.sort() |
|
|
|
for idx in range(len(input_img_list)-1): |
|
img1 = cv2.imread(os.path.join(input_file, input_img_list[idx])) |
|
img2 = cv2.imread(os.path.join(input_file, input_img_list[idx+1])) |
|
|
|
image1 = cv2.cvtColor(img1, cv2.COLOR_BGR2RGB).astype(np.float32) / np.float32(255) |
|
image2 = cv2.cvtColor(img2, cv2.COLOR_BGR2RGB).astype(np.float32) / np.float32(255) |
|
image1 = torch.from_numpy(image1).unsqueeze(0).permute(0, 3, 1, 2) |
|
image2 = torch.from_numpy(image2).unsqueeze(0).permute(0, 3, 1, 2) |
|
|
|
results = [image1, image2] |
|
|
|
inter_frames = int(inter_frames) |
|
idxes = [0, inter_frames + 1] |
|
remains = list(range(1, inter_frames + 1)) |
|
|
|
splits = torch.linspace(0, 1, inter_frames + 2) |
|
|
|
for _ in range(len(remains)): |
|
starts = splits[idxes[:-1]] |
|
ends = splits[idxes[1:]] |
|
distances = ((splits[None, remains] - starts[:, None]) / (ends[:, None] - starts[:, None]) - .5).abs() |
|
matrix = torch.argmin(distances).item() |
|
start_i, step = np.unravel_index(matrix, distances.shape) |
|
end_i = start_i + 1 |
|
|
|
x0 = results[start_i] |
|
x1 = results[end_i] |
|
|
|
x0 = x0.half() |
|
x1 = x1.half() |
|
x0 = x0.to(device) |
|
x1 = x1.to(device) |
|
|
|
dt = x0.new_full((1, 1), (splits[remains[step]] - splits[idxes[start_i]])) / (splits[idxes[end_i]] - splits[idxes[start_i]]) |
|
|
|
with torch.no_grad(): |
|
prediction = model(x0, x1, dt) |
|
insert_position = bisect.bisect_left(idxes, remains[step]) |
|
idxes.insert(insert_position, remains[step]) |
|
results.insert(insert_position, prediction.clamp(0, 1).cpu().float()) |
|
del remains[step] |
|
|
|
frames = [(tensor[0] * 255).byte().flip(0).permute(1, 2, 0).numpy().copy() for tensor in results] |
|
|
|
for sub_idx in range(len(frames)): |
|
img_path = os.path.join(image_save_dir, f'{sub_idx+idx*(inter_frames+1):06d}.png') |
|
cv2.imwrite(img_path, frames[sub_idx]) |
|
|
|
final_frames = [] |
|
final_img_list = os.listdir(image_save_dir) |
|
final_img_list.sort() |
|
for item in final_img_list: |
|
final_frames.append(cv2.imread(os.path.join(image_save_dir, item))) |
|
w, h = final_frames[0].shape[1::-1] |
|
fourcc = cv2.VideoWriter_fourcc('m', 'p', '4', 'v') |
|
video_save_dir = input_file + '.mp4' |
|
writer = cv2.VideoWriter(video_save_dir, fourcc, fps, (w, h)) |
|
for frame in final_frames: |
|
writer.write(frame) |
|
writer.release() |
|
|
|
shutil.rmtree(image_save_dir) |
|
|
|
return video_save_dir |
|
|
