Develop3

README.md

@@ -1,13 +1,24 @@
 # XFusion
 Deep learning-based spatiotemporal fusion for high-fidelity ultra-high-speed x-ray radiography  
 A model to reconstruct high quality x-ray images by combining the high spatial resolution of high-speed camera and high temporal resolution of ultra-high-speed camera image sequences.  
+Models and methods for spatial and temporal calibration are also included.
 
 ## Prerequisites
-This implementation is based on the [BasicSR toolbox](https://github.com/XPixelGroup/BasicSR). Data for model pre-training are collected from the [REDS dataset](https://seungjunnah.github.io/Datasets/reds).  
+This implementation is based on the [BasicSR toolbox](https://github.com/XPixelGroup/BasicSR) and the [RAFT model](https://github.com/princeton-vl/RAFT). Data for model pre-training are collected from the [REDS dataset](https://seungjunnah.github.io/Datasets/reds).  
 
 ## Usage
 ### Package description
-Currently, xfusion supports 2 model familities for high-quality xray image sequence reconstruction- the EDVR and Swin vision transformer, respectively.
+Currently, xfusion supports 2 model familities for high-quality xray image sequence reconstruction- the EDVR and Swin vision transformer, respectively.  
+In addition, xfusion is functional with x-ray image sequences from both the actual (data type: actual) and virtual (data type: virtual) experiments.  
+* Actual experiments utilize the dual-detector setup for data acquisition and virtual experiments synthetize dual-detector data from single-camera experiments.  
+
+With the actual dual-detector experimental data, xfusion first conducts the spatial and temporal calibration of the two image sequences.  
+* Temporal calibration is specific to the choice of the Phantom TMX7510 (Vision Research Inc., USA) camera and Shimadzu HPV-X2 (Shimadzu Corp., Japan) camera.  
+* Spatial calibration is based on the use of an electroformed copper grid (G400, Gilder Grids Ltd., UK).  
+
+Then at inference time, xfusion traverses the dual-camera image sequences in two ways: continuous and double-interval.  
+* Under the continuous mode, high-quality xray images corresponding to each Shimadzu camera image are reconstructed.  
+* Under the double-interval mode, independent high-quality xray images corresponding to each Phantom camera image are reconstructed to evaluate performance metrics.
 
 ### Package installation
 Navigate to the project root directory and then run
@@ -19,7 +30,7 @@ to install the package to the selected virtual environment.
 ### Initialization
 Run
 <pre>
-xfusion init --model_type <i>[EDVRModel or SwinIRModel]</i>
+xfusion init --model_type <i>[EDVRModel or SwinIRModel]</i> --data_type <i>[virtual or actual]</i>
 </pre>
 After initialization, a configuration file "xfusion.conf" will be generated in the home directory. This configuration file will be updated automatically within the workflow of the xfusion package.
 
@@ -38,9 +49,27 @@ xfusion convert --dir-lo-convert <i>[directory/to/low resolution/RGB/training im
 </pre>
 
 ### Training
-Run
+To train the model with data from the virtual experiments, run
+<pre>
+xfusion train --dir-lo-train <i>[directory/to/low resolution/gray-scale/training image]</i> --dir-hi-train <i>[directory/to/high resolution/gray-scale/training image]</i> --dir-lo-val <i>[directory/to/low resolution/gray-scale/validation image]</i> --dir-hi-val <i>[directory/ti/high resolution/gray-scale/validation image]</i> --opt <i>directory/to/training setting/yaml file</i> --path-train-meta-info-file <i>[directory/to/training image/meta data]</i> --path-val-meta-info-file <i>[directory/to/validation image/meta data]</i> --pretrain_network_g <i>[path/to/model weight/file/for/model initialization]</i>
+</pre>
+To train the model with data from the actual experiments, run
+<pre>
+xfusion train --dir-lo-train <i>[directory/to/low resolution/gray-scale/paired training image]</i> --dir-hi-train <i>[directory/to/high resolution/gray-scale/paired training image]</i> --dir-lo-val <i>[directory/to/low resolution/gray-scale/paired validation image]</i> --dir-hi-val <i>[directory/ti/high resolution/gray-scale/paired validation image]</i> --dataroot_context <i>[directory/to/low resolution/gray-scale/training image full]</i> --opt <i>directory/to/training setting/yaml file</i> --path-train-meta-info-file <i>[directory/to/training image/meta data]</i> --path-val-meta-info-file <i>[directory/to/validation image/meta data]</i> --meta_info_context <i>[directory/to/paired/high/lo/resolution/training image/index map]</i> --pretrain_network_g <i>[path/to/model weight/file/for/model initialization]</i>
+</pre>
+
+### Train calibration model
+To train the calibration model with data from the virtual experiments (pretrain), run
+<pre>
+xfusion train_calibration --image_root <i>[root/directory/to/training images]</i> --meta_info_file <i>[directory/to/training image/meta data]</i>
+--evaluator_weights <i>[path/to/spatiotemporal/fusion/model weight/file/for/training loss]</i>
+--flow_weights None --full_weights None
+</pre>
+To train the calibration model with data from the actual experiments (fine tune), run
 <pre>
-xfusion train --dir-lo-train <i>[directory/to/low resolution/gray-scale/training image]</i> --dir-hi-train <i>[directory/to/high resolution/gray-scale/training image]</i> --dir-lo-val <i>[directory/to/low resolution/gray-scale/validation image]</i> --dir-hi-val <i>[directory/ti/high resolution/gray-scale/validation image]</i> --opt <i>directory/to/training setting/yaml file</i> --path-train-meta-info-file <i>[directory/to/training image/meta data]</i> --path-val-meta-info-file <i>[directory/to/validation image/meta data]</i> --pretrain_network_g <i>[directory/to/model weight/file/for/model initialization]</i>
+xfusion train_calibration --image_root <i>[root/directory/to/training images]</i> --meta_info_file <i>[directory/to/training image/meta data]</i>
+--evaluator_weights <i>[path/to/spatiotemporal/fusion/model weight/file/for/training loss]</i>
+--full_weights <i>[path/to/pretrained/calibration/model weight/file/for/model initialization]</i>
 </pre>
 
 ### Test data
@@ -49,9 +78,28 @@ To download test data, run
 xfusion download --dir-inf <i>[tomobank/link/address/of/test/dataset]</i> --out-dir-inf <i>[directory/to/testing image]</i>
 </pre>
 
-### Inference
+### Spatial calibration
 Run
 <pre>
-xfusion inference --opt <i>directory/to/testing dataset/setting/yaml file</i> --arch-opt <i>directory/to/training setting/yaml file</i> --model_file <i>[path/to/model file]</i> --machine <i>tomo or polaris</i>
+xfusion calibrate --cal_model_file <i>[path/to/calibration/model file]</i> --cal_dir <i>[root/directory/to/calibration/experiment/data]</i>
+--cal_id <i>[calibration experiment subdirectory name]</i>
+</pre>
+
+### Inference
+For virtual experiment data, run
+<pre>
+xfusion inference --opt <i>directory/to/testing dataset/setting/yaml file</i> --arch-opt <i>directory/to/model architecture/yaml file</i> --img-class <i>dataset1 or dataset2</i> --model_file <i>[path/to/model file]</i> --machine <i>tomo or polaris</i>
 </pre>
-Currently to work for EDVRModel in the single-process mode and SwinIRModel in the multi-process mode.
+For actual experiment data, run
+<pre>
+xfusion inference --arch-opt <i>directory/to/model architecture/yaml file</i> --img-class <i>name of high-speed experiment</i> --model_file <i>[path/to/model file]</i> --traverse_mode <i>double or continuous</i> --tfm_file <i>[path/to/rigid/body/spatial/transformation/file]</i> --meta_file_scale <i>[path/to/nonuniform/scaling/parameters]</i> --input_dir <i>[root/directory/to/high/speed/experiment/data]</i> --case_list <i>experiment subdirectory names</i>
+</pre>
+
+### Suggested usage
+* For virtual experiment data, currently xfusion works for EDVRModel in the single-process mode and SwinIRModel in the multi-process mode.  
+When working on the virtual experiment data, first initialize the software, and then do the inference.  
+
+* For actual experiment data, xfusion works for both models in a unified way.  
+Do an independent initialization every time the model (EDVR-STF or SWIN-XVR) or the type of data (virtual or actual) changes.  
+When working on the actual experiment data, first initialize the software, then do the spatial calibration, and last do the inference.  
+Multiple experiments sharing the same spatial calibration can be grouped to run the model inference.

xfusion/config.py

@@ -20,11 +20,19 @@
            'parse_known_args',
            'write']
 def list_of_ints(arg):
+    if ',' in arg:
+        return list(map(int,arg.split(',')))
+    else:
+        if arg.lower() != 'none':
+            return [int(arg)]
+        else:
+            return None
+
+def list_of_strs(arg):
     if ',' in arg:
         return list(arg.split(','))
     else:
         return [arg]
-
 CONFIG_FILE_NAME = os.path.join(str(pathlib.Path.home()), 'xfusion.conf')
 XFUSION_HOME = os.path.join(str(pathlib.Path.home()), 'xfusion')
 XFUSION_TRAIN_HOME = os.path.join(str(XFUSION_HOME), 'train')
@@ -188,10 +196,49 @@ def list_of_ints(arg):
         'help': 'Relative frame separation in the hr images after and before the augmentation'
     },
     'poisson_noise_b0_exp': {
-        'default': [],
+        'default': 'none',
         'type': list_of_ints,
         'help': 'upper and lower bounds of the blank scan factor in log scale'
     },
+     # newly added ddp configs
+    'iter_div_factor': {
+            'default': 10,
+            'type': int,
+            'help': "Divisor of the total number of iterations"
+    },
+    'batch_size': {
+            'default': 1,
+            'type': int,
+            'help': "Number of samples in a batch"
+    },
+    'num_workers': {
+            'default': 3,
+            'type': int,
+            'help': "Number of workers for data loader"},
+    'warmup_iter': {
+        'default': -1,
+        'type': int,
+        'help': "Number of iterations for training warm up"},
+    'dataset_enlarge_ratio': {
+        'default': 10.0,
+        'type': float,
+        'help': "Multiplier of the original number of training samples"},
+    'sub_dataset': {
+        'default': False,
+        'type': lambda x: bool(strtobool(x)),
+        'help': "When set the original training data are randomly sub-sampled for training"},
+    'sub_dataset_frac': {
+        'default': 1.0,
+        'type': float,
+        'help': "Fraction of the original training data set size to sub-sample"},
+    'single_epoch_ok': {
+        'default': False,
+        'type': lambda x: bool(strtobool(x)),
+        'help': "If to only train data for one epoch"},
+    'profiler_iter_num': {
+        'default': 3,
+        'type': int,
+        'help': "Number of subsequent iterations to track computation performance for HTA"},
   }
 
 MODEL_TYPES = {'train':OrderedDict(),'inference':OrderedDict()}
@@ -225,50 +272,17 @@ def list_of_ints(arg):
             'type': int,
             'help': "Size of convolution window if used in the vit residual connection"},
         # newly added ddp configs
-        'iter_div_factor': {
-            'default': 10,
-            'type': int,
-            'help': "Divisor of the total number of iterations"},
+
         'num_frame': {
             'default': 3,
             'type': int,
             'help': "Number of input low resolution frames to the model"},
-        'batch_size': {
-            'default': 1,
-            'type': int,
-            'help': "Number of samples in a batch"},
+
         'initial_lr': {
-            'default': 2e-4,
+            'default': 5e-6,
             'type': float,
             'help': "Initial learning rate"},
-        'num_workers': {
-            'default': 3,
-            'type': int,
-            'help': "Number of workers for data loader"},
-        'warmup_iter': {
-            'default': -1,
-            'type': int,
-            'help': "Number of iterations for training warm up"},
-        'dataset_enlarge_ratio': {
-            'default': 10.0,
-            'type': float,
-            'help': "Multiplier of the original number of training samples"},
-        'sub_dataset': {
-            'default': False,
-            'type': lambda x: bool(strtobool(x)),
-            'help': "When set the original training data are randomly sub-sampled for training"},
-        'sub_dataset_frac': {
-            'default': 1.0,
-            'type': float,
-            'help': "Fraction of the original training data set size to sub-sample"},
-        'single_epoch_ok': {
-            'default': False,
-            'type': lambda x: bool(strtobool(x)),
-            'help': "If to only train data for one epoch"},
-        'profiler_iter_num': {
-            'default': 3,
-            'type': int,
-            'help': "Number of subsequent iterations to track computation performance for HTA"},}
+        }
 
 MODEL_TYPES['train']['EDVRModel'] = {
     'initial_lr': {
@@ -326,7 +340,7 @@ def list_of_ints(arg):
         'default': 'none',
         'type': Path,
         'help': "When set the path to the model weights to initialize the model"},
-    'num_frame': {
+    'num_frame_infer': {
         'default': 3,
         'type': int,
         'help': "Number of input low resolution frames to the model"},
@@ -510,7 +524,7 @@ def list_of_ints(arg):
         'help': 'when set the model has separate encoders for the hr/lr input image tensors'
     },
     'cal_train_poisson_b0_exponent': {
-        'default': [],
+        'default': 'none',
         'type': list_of_ints,
         'help': 'upper and lower bounds of the blank scan factor in log scale'
     },
@@ -688,7 +702,7 @@ def list_of_ints(arg):
         'help': "Full path to the user experiment data"
     },
     'case_list': {'default': [],
-        'type': list_of_ints,
+        'type': list_of_strs,
         'help': "Names of the individual experiment subdirectories"
     },
     'device': {'default': 0,

xfusion/inference/infer_swin_ddp.py

@@ -59,10 +59,10 @@ def run_inference(ts, args):
 
     opt['network_g']['embed_dim'] = args.embed_dim
     opt['network_g']['num_feat_ext'] = args.embed_dim
-    opt['network_g']['num_frame'] = args.num_frame
+    opt['network_g']['num_frame'] = args.num_frame_infer
 
     opt['network_g']['depths'] = [d * args.scale_depth for d in opt['network_g']['depths']]
-    opt['network_g']['window_size'][0] += (args.num_frame - 3)
+    opt['network_g']['window_size'][0] += (args.num_frame_infer - 3)
     opt['network_g']['window_size'][1:] = [args.window_size_spatial] * len(opt['network_g']['window_size'][1:])
 
     opt['network_g']['num_heads'] = [int(np.ceil(h * args.scale_mha)) for h in opt['network_g']['num_heads']]
@@ -75,7 +75,7 @@ def run_inference(ts, args):
     dataroot = config.get_inf_data_dirs(img_class)
     inf_home_dir = Path(dataroot).parent
 
-    out_dir = Path(inf_home_dir / f'{test_set_name}_{mode}_align_{args.align_features_ok}_prepostx_fixed_lr_r_{lo_frame_sep}_hr_d_{hi_frame_sep*2}_b0_{b0}_em_{args.embed_dim}_nf_{args.num_frame}_dep_{args.scale_depth}_nh_{args.scale_mha}_win_{args.window_size_spatial}_model_{args.model_file.parent.stem}_ddp_swinVRHP_rem1_{ts}')
+    out_dir = Path(inf_home_dir / f'{test_set_name}_{mode}_align_{args.align_features_ok}_prepostx_fixed_lr_r_{lo_frame_sep}_hr_d_{hi_frame_sep*2}_b0_{b0}_em_{args.embed_dim}_nf_{args.num_frame_infer}_dep_{args.scale_depth}_nh_{args.scale_mha}_win_{args.window_size_spatial}_model_{args.model_file.parent.stem}_ddp_swinVRHP_rem1_{ts}')
     out_dir.mkdir(exist_ok=True,parents=True)
 
     log_file = os.path.join(out_dir, f"inference.log")
@@ -129,7 +129,7 @@ def run_inference(ts, args):
         #    continue
         dataset_opt['scale'] = 4
         dataset_opt['gt_size'] = gt_size[:2]
-        dataset_opt['num_frame'] = args.num_frame
+        dataset_opt['num_frame'] = args.num_frame_infer
         #test_set = VideoTestDatasetSTF(dataset_opt)
         #test_set = build_dataset(dataset_opt)
         test_set = XrayVideoTestDatasetSTF(dataset_opt)

xfusion/inference/infer_with_auto_calibration_main4_continuous2.py

@@ -79,7 +79,7 @@ def run_inference(args):
     exp_root = Path(config.get_inf_data_dirs(input_dir.name))
     device = args.device
 
-    opt_path = args.opt
+    opt_path = args.arch_opt
     opt = yaml_load(opt_path)
     opt['manual_seed'] = 10
     torch.manual_seed(opt['manual_seed'])

xfusion/train/basicsr/metrics/metric_util.py

@@ -15,8 +15,6 @@
 
 import numpy as np
 
-from xfusion.train.basicsr.utils import bgr2ycbcr
-
 
 def reorder_image(img, input_order='HWC'):
     """Reorder images to 'HWC' order.
@@ -42,19 +40,3 @@ def reorder_image(img, input_order='HWC'):
     if input_order == 'CHW':
         img = img.transpose(1, 2, 0)
     return img
-
-
-def to_y_channel(img):
-    """Change to Y channel of YCbCr.
-
-    Args:
-        img (ndarray): Images with range [0, 255].
-
-    Returns:
-        (ndarray): Images with range [0, 255] (float type) without round.
-    """
-    img = img.astype(np.float32) / 255.
-    if img.ndim == 3 and img.shape[2] == 3:
-        img = bgr2ycbcr(img, y_only=True)
-        img = img[..., None]
-    return img * 255.

xfusion/train/basicsr/metrics/niqe.py

@@ -20,7 +20,7 @@
 from scipy.ndimage import convolve
 from scipy.special import gamma
 
-from xfusion.train.basicsr.metrics.metric_util import reorder_image, to_y_channel
+from xfusion.train.basicsr.metrics.metric_util import reorder_image
 from xfusion.train.basicsr.utils.matlab_functions import imresize
 from xfusion.train.basicsr.utils.registry import METRIC_REGISTRY
 
@@ -198,7 +198,7 @@ def calculate_niqe(img, crop_border, input_order='HWC', convert_to='y', **kwargs
     if input_order != 'HW':
         img = reorder_image(img, input_order=input_order)
         if convert_to == 'y':
-            img = to_y_channel(img)
+            raise Exception('not implemented yet...')
         elif convert_to == 'gray':
             img = cv2.cvtColor(img / 255., cv2.COLOR_BGR2GRAY) * 255.
         img = np.squeeze(img)

xfusion/train/basicsr/metrics/psnr_ssim.py

@@ -18,8 +18,7 @@
 import torch
 import torch.nn.functional as F
 
-from xfusion.train.basicsr.metrics.metric_util import reorder_image, to_y_channel
-from xfusion.train.basicsr.utils.color_util import rgb2ycbcr_pt
+from xfusion.train.basicsr.metrics.metric_util import reorder_image
 from xfusion.train.basicsr.utils.registry import METRIC_REGISTRY
 
 
@@ -51,8 +50,7 @@ def calculate_psnr(img, img2, crop_border, input_order='HWC', test_y_channel=Fal
         img2 = img2[crop_border:-crop_border, crop_border:-crop_border, ...]
 
     if test_y_channel:
-        img = to_y_channel(img)
-        img2 = to_y_channel(img2)
+        raise Exception('not implemented yet...')
 
     img = img.astype(np.float64)
     img2 = img2.astype(np.float64)
@@ -86,8 +84,7 @@ def calculate_psnr_pt(img, img2, crop_border, test_y_channel=False, **kwargs):
         img2 = img2[:, :, crop_border:-crop_border, crop_border:-crop_border]
 
     if test_y_channel:
-        img = rgb2ycbcr_pt(img, y_only=True)
-        img2 = rgb2ycbcr_pt(img2, y_only=True)
+        raise Exception('not implemented yet...')
 
     img = img.to(torch.float64)
     img2 = img2.to(torch.float64)
@@ -131,8 +128,7 @@ def calculate_ssim(img, img2, crop_border, input_order='HWC', test_y_channel=Fal
         img2 = img2[crop_border:-crop_border, crop_border:-crop_border, ...]
 
     if test_y_channel:
-        img = to_y_channel(img)
-        img2 = to_y_channel(img2)
+        raise Exception('not implemented yet...')
 
     img = img.astype(np.float64)
     img2 = img2.astype(np.float64)
@@ -172,8 +168,7 @@ def calculate_ssim_pt(img, img2, crop_border, test_y_channel=False, **kwargs):
         img2 = img2[:, :, crop_border:-crop_border, crop_border:-crop_border]
 
     if test_y_channel:
-        img = rgb2ycbcr_pt(img, y_only=True)
-        img2 = rgb2ycbcr_pt(img2, y_only=True)
+        raise Exception('not implemented yet...')
 
     img = img.to(torch.float64)
     img2 = img2.to(torch.float64)