Add 32x inference

.gitignore

@@ -3,3 +3,4 @@
 *.py#
 
 *.pkl
+output/*

inference.py

@@ -1,3 +1,4 @@
+import os
 import cv2
 import torch
 import argparse
@@ -8,21 +9,37 @@ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
 parser = argparse.ArgumentParser(description='Interpolation for a pair of images')
 parser.add_argument('--img', dest='img', nargs=2, required=True)
+parser.add_argument('--times', default=5, type=int)
 args = parser.parse_args()
 
 model = Model()
 model.load_model('./train_log')
 model.eval()
 model.device()
+    
 img0 = cv2.imread(args.img[0])
 img1 = cv2.imread(args.img[1])
-h, w, _ = img0.shape
+
+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)
+n, c, h, w = img0.shape
 ph = h // 32 * 32
 pw = w // 32 * 32
 padding = (0, pw - w, 0, ph - h)
-img0 = torch.tensor(img0.transpose(2, 0, 1)).to(device) / 255.
-img1 = torch.tensor(img1.transpose(2, 0, 1)).to(device) / 255.
-imgs = F.pad(torch.cat((img0, img1), 0).float(), padding)
-with torch.no_grad():
-    res = model.inference(imgs.unsqueeze(0)) * 255
-cv2.imwrite('output.png', res[0].numpy().transpose(1, 2, 0)[:h, :w])
+img0 = F.pad(img0, padding)
+img1 = F.pad(img1, padding)
+
+img_list = [img0, img1]
+for i in range(args.times):
+    tmp = []
+    for j in range(len(img_list) - 1):
+        mid = model.inference(img_list[j], img_list[j + 1])
+        tmp.append(img_list[j])
+        tmp.append(mid)
+    tmp.append(img1)
+    img_list = tmp
+
+if not os.path.exists('output'):
+    os.mkdir('output')
+for i in range(len(img_list)):
+    cv2.imwrite('output/img{}.png'.format(i), img_list[i][0].numpy().transpose(1, 2, 0)[:h, :w] * 255)

model/RIFE.py

@@ -200,10 +200,11 @@ class Model:
         else:
             return pred
 
-    def inference(self, imgs):
+    def inference(self, img0, img1):
         with torch.no_grad():
+            imgs = torch.cat((img0, img1), 1)
             flow, _ = self.flownet(imgs)
-        return self.predict(imgs, flow, training=False)
+        return self.predict(imgs, flow, training=False).detach()
 
     def update(self, imgs, gt, learning_rate=0, mul=1, training=True, flow_gt=None):
         for param_group in self.optimG.param_groups: