FEAT: Adding float32 support for RIFE

Author: Cabana-HPC1

tomosuitepy/base/reconstruct.py

@@ -504,7 +504,10 @@ def prepare_rife(basedir, start_row, end_row, rife_types, verbose):
 
     for file in tqdm(new_files, desc='Loading Data'):
         original = cv2.imread(file, -1)
-        fixed_original = rgb2gray(original)
+        if len(np.shape(original)) == 3:
+            fixed_original = rgb2gray(original)
+        else:
+            fixed_original = original
         if rife_types[2] == True:
             fixed_original = np.log(fixed_original)
         #fixed_original *= 255.0
@@ -515,7 +518,7 @@ def prepare_rife(basedir, start_row, end_row, rife_types, verbose):
     shape = prj_data.shape[0]
 
     _theta = np.load(f"{basedir}extracted/theta/theta.npy")
-    theta = np.linspace(_theta[0], theta[-1], prj_data.shape[0])
+    theta = np.linspace(_theta[0], _theta[-1], prj_data.shape[0])
 
     if verbose:
         print(f"The shape of this data is: {prj_data.shape}")

tomosuitepy/easy_networks/rife/data_prep.py

@@ -62,6 +62,7 @@ def deal_with_sparse_angle(prj_data, theta,
 
     return prj_data, theta
 
+
 def view_prj_contrast(basedir, cutoff=None, above_or_below='below',
                         analysis_func=np.sum, plot=True):
 
@@ -110,7 +111,6 @@ def view_prj_contrast(basedir, cutoff=None, above_or_below='below',
     return np.asarray(analysis_idx), prj_data[analysis_idx], theta[analysis_idx]
 
 
-
 def create_prj_mp4(basedir, video_type='input', types='base', force_positive=False,
                    sparse_angle_removal=0, fps=30, torf=False, apply_exp=False, prj_idx=None):
     """
@@ -213,6 +213,7 @@ def create_prj_mp4(basedir, video_type='input', types='base', force_positive=Fal
     return prj_data, np.asarray(out_data)
 
 
+
 def rife_predict(basedir, location_of_rife=rife_path, exp=2, scale=1.0,
                  gpu='0', video_input_type='input',
                  video_output_type='predicted',
@@ -348,4 +349,42 @@ def create_prj_mp4_old(basedir, output_file, types='base', sparse_angle_removal=
         im *= 255.0
         out.write((im).astype(np.uint8))
 
-    out.release()
+    out.release()
+    
+
+def full_res_rife(basedir, location_of_rife=rife_path, exp=2,
+                 gpu='0', output_folder='frames', python_location=''):
+    """
+    Use the neural network called RIFE to upscale the amount of projections.
+
+    Parameters
+    ----------
+    basedir : str
+        Path to the project.
+
+    location_of_rife : str
+        Path to the github repo of RIFE with / at the end.
+
+    exp : int
+        2 to the power of exp that the frames will be upscaled by
+
+    gpu : str
+        The string index of the gpu to use.
+
+    output_folder : str
+        The name of the output folder to be created at {basedir}rife/{output_folder}/
+
+    Returns
+    -------
+    command
+        A command to be used in a terminal with the RIFE conda env variables installed.
+    """
+
+    pre = f'cd {location_of_rife} &&'
+    first = f'{python_location}python inference_img.py'
+    second = f'--exp={exp}'
+    third = f'--basedir={basedir}'
+    fourth_inter = f'--gpu={gpu}'
+    fifth = f"--output={output_folder}"
+
+    return f"{pre} {first} {second} {third} {fourth_inter} {fifth}"

tomosuitepy/hard_networks/RIFE/arXiv2020-RIFE/inference_img.py

@@ -1,10 +1,29 @@
-import os
-import cv2
+import os, sys, warnings, cv2, argparse
+
+parser = argparse.ArgumentParser(description='Interpolation for a pair of images')
+parser.add_argument('--basedir', type=str, required=True)
+parser.add_argument('--exp', default=2, type=int)
+parser.add_argument('--output', default='frames_check', type=str)
+parser.add_argument('--gpu', default=0, type=int)
+
+args = parser.parse_args()
+args.ratio = float(0.0) # inference ratio between two images with 0 - 1 range
+args.rthreshold = float(0.02) # returns image when actual ratio falls in given range threshold
+args.rmaxcycles = int(8) # limit max number of bisectional cycles
+
+os.environ['CUDA_VISIBLE_DEVICES'] = str(args.gpu)
+
 import torch
-import argparse
 from torch.nn import functional as F
 from model.RIFE_HDv2 import Model
-import warnings
+import tifffile as tif
+import numpy as np
+from tqdm import tqdm
+from skimage.color import rgb2gray
+
+sys.path.append('../../..')
+from base.common import load_extracted_prj
+
 warnings.filterwarnings("ignore")
 
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
@@ -13,78 +32,79 @@
     torch.backends.cudnn.enabled = True
     torch.backends.cudnn.benchmark = True
 
-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')
-args = parser.parse_args()
-
 model = Model()
 model.load_model(os.path.join(os.path.dirname(os.path.realpath(__file__)), 'train_log'), -1)
 model.eval()
 model.device()
 
-if args.img[0].endswith('.exr') and args.img[1].endswith('.exr'):
-    img0 = cv2.imread(args.img[0], cv2.IMREAD_COLOR | cv2.IMREAD_ANYDEPTH)
-    img1 = cv2.imread(args.img[1], 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)
+# Load in all the images
+total_prjs = load_extracted_prj(args.basedir)
+save_location = f"{args.basedir}rife/{args.output}/"
+prj_max = total_prjs.max()
+total_prjs = total_prjs / prj_max
+total_prjs = total_prjs * 255.0
 
-else:
-    img0 = cv2.imread(args.img[0])
-    img1 = cv2.imread(args.img[1])
-    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)
+# Iterate over the images
+zfill_val = len(str(len(total_prjs) * 2**args.exp))
+current_frame = 0
 
-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)
+for iteration in tqdm(range(0, len(total_prjs)-1), desc='Interpolation'):
+    # Tripple the image arrays
+    img0 = np.dstack((total_prjs[iteration], total_prjs[iteration], total_prjs[iteration]))
+    img1 = np.dstack((total_prjs[iteration + 1], total_prjs[iteration + 1], total_prjs[iteration + 1]))
 
+    img0 = img0.astype(np.float32)
+    img1 = img1.astype(np.float32)
 
-if args.ratio:
-    img_list = [img0]
-    img0_ratio = 0.0
-    img1_ratio = 1.0
-    if args.ratio <= img0_ratio + args.rthreshold / 2:
-        middle = img0
-    elif args.ratio >= img1_ratio - args.rthreshold / 2:
-        middle = img1
-    else:
-        tmp_img0 = img0
-        tmp_img1 = img1
-        for inference_cycle in range(args.rmaxcycles):
-            middle = model.inference(tmp_img0, tmp_img1)
-            middle_ratio = ( img0_ratio + img1_ratio ) / 2
-            if args.ratio - (args.rthreshold / 2) <= middle_ratio <= args.ratio + (args.rthreshold / 2):
-                break
-            if args.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(args.exp):
-        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
+    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 - 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 not os.path.exists('output'):
-    os.mkdir('output')
-for i in range(len(img_list)):
-    if args.img[0].endswith('.exr') and args.img[1].endswith('.exr'):
-        cv2.imwrite('output/img{}.exr'.format(i), (img_list[i][0]).cpu().numpy().transpose(1, 2, 0)[:h, :w], [cv2.IMWRITE_EXR_TYPE, cv2.IMWRITE_EXR_TYPE_HALF])
+    if args.ratio:
+        img_list = [img0]
+        img0_ratio = 0.0
+        img1_ratio = 1.0
+        if args.ratio <= img0_ratio + args.rthreshold / 2:
+            middle = img0
+        elif args.ratio >= img1_ratio - args.rthreshold / 2:
+            middle = img1
+        else:
+            tmp_img0 = img0
+            tmp_img1 = img1
+            for inference_cycle in range(args.rmaxcycles):
+                middle = model.inference(tmp_img0, tmp_img1)
+                middle_ratio = ( img0_ratio + img1_ratio ) / 2
+                if args.ratio - (args.rthreshold / 2) <= middle_ratio <= args.ratio + (args.rthreshold / 2):
+                    break
+                if args.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:
-        cv2.imwrite('output/img{}.png'.format(i), (img_list[i][0] * 255).byte().cpu().numpy().transpose(1, 2, 0)[:h, :w])
+        img_list = [img0, img1]
+        for i in range(args.exp):
+            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(save_location):
+        os.mkdir(save_location)
+    for i in range(len(img_list)):
+        im2save = (img_list[i][0] * 255).cpu().numpy().transpose(1, 2, 0)[:h, :w]
+        im2save = rgb2gray(im2save)
+        tif.imsave(f'{save_location}/img_{str(current_frame).zfill(zfill_val)}.tif', im2save)
+        current_frame += 1