Spaces:
Sleeping
Sleeping
| import cv2 | |
| import torch | |
| from torch.nn import functional as F | |
| import warnings | |
| import os | |
| warnings.filterwarnings("ignore") | |
| device = torch.device("cuda" if torch.cuda.is_available() else "cpu") | |
| torch.set_grad_enabled(False) | |
| if torch.cuda.is_available(): | |
| torch.backends.cudnn.enabled = True | |
| torch.backends.cudnn.benchmark = True | |
| class RIFE: | |
| _instance = None | |
| def __new__(cls, model_path='./ckpt_models/rife'): | |
| if cls._instance is None: | |
| cls._instance = super(RIFE, cls).__new__(cls) | |
| cls._instance.initialize(model_path) | |
| return cls._instance | |
| def initialize(self, model_path): | |
| try: | |
| try: | |
| from model.RIFE_HDv2 import Model | |
| self.model = Model() | |
| self.model.load_model(model_path, -1) | |
| print("Loaded v2.x HD model.") | |
| except: | |
| from train_log.RIFE_HDv3 import Model | |
| self.model = Model() | |
| self.model.load_model(model_path, -1) | |
| print("Loaded v3.x HD model.") | |
| except: | |
| try: | |
| from model.RIFE_HD import Model | |
| self.model = Model() | |
| self.model.load_model(model_path, -1) | |
| print("Loaded v1.x HD model") | |
| except: | |
| from model.RIFE import Model | |
| self.model = Model() | |
| self.model.load_model(model_path, -1) | |
| print("Loaded ArXiv-RIFE model") | |
| self.model.eval() | |
| self.model.device() | |
| def interpolate(self, img0, img1, exp=4, ratio=0, rthreshold=0.02, rmaxcycles=8): | |
| if isinstance(img0, str) and isinstance(img1, str): | |
| if img0.endswith('.exr') and img1.endswith('.exr'): | |
| img0 = cv2.imread(img0, cv2.IMREAD_COLOR | cv2.IMREAD_ANYDEPTH) | |
| img1 = cv2.imread(img1, cv2.IMREAD_COLOR | cv2.IMREAD_ANYDEPTH) | |
| img0 = (torch.tensor(img0.transpose(2, 0, 1)).to(device)).unsqueeze(0) | |
| img1 = (torch.tensor(img1.transpose(2, 0, 1)).to(device)).unsqueeze(0) | |
| else: | |
| img0 = cv2.imread(img0, cv2.IMREAD_UNCHANGED) | |
| img1 = cv2.imread(img1, cv2.IMREAD_UNCHANGED) | |
| img0 = (torch.tensor(img0.transpose(2, 0, 1)).to(device) / 255.).unsqueeze(0) | |
| img1 = (torch.tensor(img1.transpose(2, 0, 1)).to(device) / 255.).unsqueeze(0) | |
| elif isinstance(img0, torch.Tensor) and isinstance(img1, torch.Tensor): | |
| img0 = img0.to(device) | |
| img1 = img1.to(device) | |
| else: | |
| raise ValueError("Input images must be either file paths or torch tensors") | |
| n, c, h, w = img0.shape | |
| ph = ((h - 1) // 32 + 1) * 32 | |
| pw = ((w - 1) // 32 + 1) * 32 | |
| padding = (0, pw - w, 0, ph - h) | |
| img0 = F.pad(img0, padding) | |
| img1 = F.pad(img1, padding) | |
| if ratio: | |
| img_list = [img0] | |
| img0_ratio = 0.0 | |
| img1_ratio = 1.0 | |
| if ratio <= img0_ratio + rthreshold / 2: | |
| middle = img0 | |
| elif ratio >= img1_ratio - rthreshold / 2: | |
| middle = img1 | |
| else: | |
| tmp_img0 = img0 | |
| tmp_img1 = img1 | |
| for inference_cycle in range(rmaxcycles): | |
| middle = self.model.inference(tmp_img0, tmp_img1) | |
| middle_ratio = (img0_ratio + img1_ratio) / 2 | |
| if ratio - (rthreshold / 2) <= middle_ratio <= ratio + (rthreshold / 2): | |
| break | |
| if ratio > middle_ratio: | |
| tmp_img0 = middle | |
| img0_ratio = middle_ratio | |
| else: | |
| tmp_img1 = middle | |
| img1_ratio = middle_ratio | |
| img_list.append(middle) | |
| img_list.append(img1) | |
| else: | |
| img_list = [img0, img1] | |
| for i in range(exp): | |
| tmp = [] | |
| for j in range(len(img_list) - 1): | |
| mid = self.model.inference(img_list[j], img_list[j + 1]) | |
| tmp.append(img_list[j]) | |
| tmp.append(mid) | |
| tmp.append(img1) | |
| img_list = tmp | |
| return [img[:, :, :h, :w] for img in img_list] | |
| def unload(self): | |
| if hasattr(self, 'model'): | |
| del self.model | |
| torch.cuda.empty_cache() | |
| print("RIFE model unloaded and CUDA cache cleared.") | |
| def reset(cls): | |
| if cls._instance is not None: | |
| cls._instance.unload() | |
| cls._instance = None | |
| print("RIFE instance reset.") | |
| def save_images(img_list, output_dir='output', img0_path='', img1_path=''): | |
| if not os.path.exists(output_dir): | |
| os.makedirs(output_dir) | |
| for i, img in enumerate(img_list): | |
| if img0_path.endswith('.exr') and img1_path.endswith('.exr'): | |
| cv2.imwrite(os.path.join(output_dir, f'img{i}.exr'), img[0].cpu().numpy().transpose(1, 2, 0), [cv2.IMWRITE_EXR_TYPE, cv2.IMWRITE_EXR_TYPE_HALF]) | |
| else: | |
| cv2.imwrite(os.path.join(output_dir, f'img{i}.png'), (img[0] * 255).byte().cpu().numpy().transpose(1, 2, 0)) | |
| if __name__ == "__main__": | |
| import argparse | |
| parser = argparse.ArgumentParser(description='Interpolation for a pair of images') | |
| parser.add_argument('--img', dest='img', nargs=2, required=True) | |
| parser.add_argument('--exp', default=4, type=int) | |
| parser.add_argument('--ratio', default=0, type=float, help='inference ratio between two images with 0 - 1 range') | |
| parser.add_argument('--rthreshold', default=0.02, type=float, help='returns image when actual ratio falls in given range threshold') | |
| parser.add_argument('--rmaxcycles', default=8, type=int, help='limit max number of bisectional cycles') | |
| parser.add_argument('--model', dest='modelDir', type=str, default='./ckpt_models/rife', help='directory with trained model files') | |
| args = parser.parse_args() | |
| rife = RIFE(args.modelDir) | |
| img_list = rife.interpolate(args.img[0], args.img[1], args.exp, args.ratio, args.rthreshold, args.rmaxcycles) | |
| save_images(img_list, img0_path=args.img[0], img1_path=args.img[1]) | |
| rife.unload() | |