reformatting codes

dataset/DeepPhysDataset.py

@@ -1,7 +1,7 @@
+import numpy as np
 import torch
-from torch.utils.data import Dataset
 import torchvision.transforms as transforms
-import numpy as np
+from torch.utils.data import Dataset
 
 
 class DeepPhysDataset(Dataset):
@@ -19,7 +19,7 @@ def __getitem__(self, index):
         motion_data = torch.tensor(np.transpose(self.m[index], (2, 0, 1)), dtype=torch.float32)
         target = torch.tensor(self.label[index], dtype=torch.float32)
 
-        inputs = [appearance_data,motion_data]
+        inputs = [appearance_data, motion_data]
 
         if torch.cuda.is_available():
             inputs.to('cuda')

dataset/PhysNetDataset.py

@@ -1,11 +1,11 @@
+import numpy as np
 import torch
-from torch.utils.data import Dataset
 import torchvision.transforms as transforms
-import numpy as np
+from torch.utils.data import Dataset
 
 
 class PhysNetDataset(Dataset):
-    def __init__(self,video_data, label_data):
+    def __init__(self, video_data, label_data):
         self.transform = transforms.Compose([transforms.ToTensor()])
         self.video_data = video_data
         self.label = label_data

log.py

@@ -6,4 +6,4 @@ def log_info_time(message, time):
 
 
 def log_warning(message):
-    print(Fore.LIGHTRED_EX + Style.BRIGHT + message + Style.RESET_ALL)
+    print(Fore.LIGHTRED_EX + Style.BRIGHT + message + Style.RESET_ALL)

main.py

@@ -8,9 +8,9 @@
 from tqdm import tqdm
 
 from dataset.dataset_loader import dataset_loader
-from log import log_info_time, log_warning
+from log import log_info_time
 from loss import loss_fn
-from models import is_model_support, get_model
+from models import is_model_support, get_model, summary
 from optim import optimizer
 from utils.dataset_preprocess import preprocessing
 from utils.funcs import normalize, plot_graph
@@ -19,6 +19,7 @@
     jsonObject = json.load(f)
     __PREPROCESSING__ = jsonObject.get("__PREPROCESSING__")
     __TIME__ = jsonObject.get("__TIME__")
+    __MODEL_SUMMARY__ = jsonObject.get("__MODEL_SUMMARY__")
     options = jsonObject.get("options")
     params = jsonObject.get("params")
     hyper_params = jsonObject.get("hyper_params")
@@ -27,7 +28,7 @@
 """
 Check Model Support
 """
-is_model_support(model_params["name"],model_params["name_comment"])
+is_model_support(model_params["name"], model_params["name_comment"])
 '''
 Generate preprocessed data hpy file 
 '''
@@ -101,6 +102,10 @@
     model.cuda()
 else:
     model = model.to('cpu')
+
+if __MODEL_SUMMARY__:
+    summary(model)
+
 if __TIME__:
     log_info_time("model initialize time \t: ", datetime.timedelta(seconds=time.time() - start_time))
 
@@ -191,6 +196,6 @@
                     else:
                         inference_array.extend(outputs.numpy())
                         target_array.extend(target.numpy())
-            plot_graph(0, len(inference_array),target_array,inference_array)
+            plot_graph(0, len(inference_array), target_array, inference_array)
         if __TIME__:
             log_info_time("inference time \t: ", datetime.timedelta(seconds=time.time() - start_time))

models.py

@@ -1,3 +1,5 @@
+import torchsummary
+
 from log import log_warning
 from nets.models.DeepPhys import DeepPhys
 from nets.models.DeepPhys_DA import DeepPhys_DA
@@ -6,7 +8,7 @@
 
 def get_model(model_name: str = "DeepPhys"):
     """
-    :param model_name: implemented model name
+    :param model_name: model name
     :return: model
     """
     if model_name == "DeepPhys":
@@ -21,6 +23,26 @@ def get_model(model_name: str = "DeepPhys"):
 
 
 def is_model_support(model_name, model_list):
+    """
+    :param model_name: model name
+    :param model_list: implemented model list
+    :return: model
+    """
     if not (model_name in model_list):
         log_warning("use implemented model")
         raise NotImplementedError("implement a custom model(%s) in /nets/models/" % model_name)
+
+
+def summary(model, model_name):
+    """
+    :param model: torch.nn.module class
+    :param model_name: implemented model name
+    :return: model
+    """
+    if model_name == "DeepPhys" or model_name == DeepPhys_DA:
+        torchsummary.summary(model, (2, 3, 36, 36))
+    elif model_name == "PhsNet":
+        torchsummary.summary(model, (3, 32, 128, 128))
+    else:
+        log_warning("use implemented model")
+        raise NotImplementedError("implement a custom model(%s) in /nets/models/" % model_name)

nets/blocks/complexLayers.py

@@ -7,12 +7,13 @@
 """
 
 import torch
-from torch.nn import Module, Parameter, init
 from torch.nn import Conv2d, Linear, BatchNorm1d, BatchNorm2d
 from torch.nn import ConvTranspose2d
-from nets.funcs.complexFunctions import complex_relu, complex_max_pool2d, complex_avg_pool2d, complex_sigmoid
+from torch.nn import Module, Parameter, init
+
 from nets.funcs.complexFunctions import complex_dropout, complex_dropout2d, complex_adaptive_avg_pool2d, \
     complex_adaptive_max_pool2d
+from nets.funcs.complexFunctions import complex_relu, complex_max_pool2d, complex_avg_pool2d, complex_sigmoid
 
 
 def apply_complex(fr, fi, input, dtype=torch.complex64):

nets/blocks/decoder_blocks.py

@@ -11,5 +11,5 @@ def __init__(self):
             DeConvBlock3D(64, 64, [4, 1, 1], [2, 1, 1], [1, 0, 0])
         )
 
-    def forward(self,x):
-        return self.decoder_block(x)
+    def forward(self, x):
+        return self.decoder_block(x)

nets/blocks/encoder_blocks.py

@@ -7,7 +7,7 @@ class encoder_block(torch.nn.Module):
     def __init__(self):
         super(encoder_block, self).__init__()
         self.encoder_block = torch.nn.Sequential(
-            ConvBlock3D(3, 16, [1, 5, 5], [1, 1, 1], [0, 2, 2]) ,
+            ConvBlock3D(3, 16, [1, 5, 5], [1, 1, 1], [0, 2, 2]),
             torch.nn.MaxPool3d((1, 2, 2), stride=(1, 2, 2)),
             ConvBlock3D(16, 32, [3, 3, 3], [1, 1, 1], [1, 1, 1]),
             ConvBlock3D(32, 64, [3, 3, 3], [1, 1, 1], [1, 1, 1]),
@@ -35,5 +35,5 @@ def __init__(self):
         # self.conv_block8 = ConvBlock3D(64, 64, [3, 3, 3], [1, 1, 1], [1, 1, 1])
         # self.conv_block9 = ConvBlock3D(64, 64, [3, 3, 3], [1, 1, 1], [1, 1, 1])
 
-    def forward(self,x):
-        return self.encoder_block(x)
+    def forward(self, x):
+        return self.encoder_block(x)

nets/blocks/motionBlocks.py

@@ -1,5 +1,6 @@
 import torch
 import torch.nn as nn
+
 from blocks import TSM_Block
 
 

nets/funcs/complexFunctions.py

@@ -5,10 +5,10 @@
 @author: spopoff
 """
 
-from torch.nn.functional import relu, max_pool2d, avg_pool2d, dropout, dropout2d, interpolate,adaptive_max_pool2d,\
-    adaptive_avg_pool2d
-from torch import sigmoid,tanh
 import torch
+from torch import sigmoid, tanh
+from torch.nn.functional import relu, max_pool2d, avg_pool2d, dropout, dropout2d, interpolate, adaptive_max_pool2d, \
+    adaptive_avg_pool2d
 
 
 def complex_matmul(A, B):
@@ -42,25 +42,33 @@ def complex_normalize(input):
 
     return real_norm.type(torch.complex64) + 1j * imag_norm.type(torch.complex64)
 
+
 def complex_mean(input):
     real_value, imag_value = input.real, input.imag
     real_mean = torch.mean(real_value, dim=1, keepdim=True)
     imag_mean = torch.mean(imag_value, dim=1, keepdim=True)
     return real_mean.type(torch.complex64) + 1j * imag_mean.type(torch.complex64)
 
+
 def complex_max(input):
     real_value, imag_value = input.real, input.imag
     real_mean, _ = torch.max(real_value, dim=1, keepdim=True)
     imag_mean, _ = torch.max(imag_value, dim=1, keepdim=True)
     return real_mean.type(torch.complex64) + 1j * imag_mean.type(torch.complex64)
 
+
 def complex_relu(input):
     return relu(input.real).type(torch.complex64) + 1j * relu(input.imag).type(torch.complex64)
+
+
 def complex_tanh(input):
     return tanh(input.real).type(torch.complex64) + 1j * tanh(input.imag).type(torch.complex64)
+
+
 def complex_sigmoid(input):
     return sigmoid(input.real).type(torch.complex64) + 1j * sigmoid(input.imag).type(torch.complex64)
 
+
 def _retrieve_elements_from_indices(tensor, indices):
     flattened_tensor = tensor.flatten(start_dim=-2)
     output = flattened_tensor.gather(dim=-1, index=indices.flatten(start_dim=-2)).view_as(indices)
@@ -118,6 +126,8 @@ def complex_max_pool2d(input, kernel_size, stride=None, padding=0,
     angle = _retrieve_elements_from_indices(angle, indices)
     return absolute_value \
            * (torch.cos(angle).type(torch.complex64) + 1j * torch.sin(angle).type(torch.complex64))
+
+
 def complex_adaptive_max_pool2d(input, output_size):
     '''
     Perform complex max pooling by selecting on the absolute value on the complex values.
@@ -138,7 +148,6 @@ def complex_adaptive_max_pool2d(input, output_size):
            * (torch.cos(angle).type(torch.complex64) + 1j * torch.sin(angle).type(torch.complex64))
 
 
-
 def complex_dropout(input, p=0.5, training=True):
     # need to have the same dropout mask for real and imaginary part,
     # this not a clean solution!
@@ -157,7 +166,8 @@ def complex_dropout2d(input, p=0.5, training=True):
     mask.type(input.dtype)
     return mask * input
 
+
 def complex_adaptive_avg_pool2d(input, *args, **kwargs):
     out_value_real = adaptive_avg_pool2d(input.real, *args, **kwargs)
     out_value_imag = adaptive_avg_pool2d(input.imag, *args, **kwargs)
-    return out_value_real.type(torch.complex64)+1j*out_value_imag.type(torch.complex64)
+    return out_value_real.type(torch.complex64) + 1j * out_value_imag.type(torch.complex64)

nets/layers/complexLayers.py

@@ -7,12 +7,13 @@
 """
 
 import torch
-from torch.nn import Module, Parameter, init
 from torch.nn import Conv2d, Linear, BatchNorm1d, BatchNorm2d
 from torch.nn import ConvTranspose2d
-from nets.funcs.complexFunctions import complex_relu, complex_max_pool2d, complex_avg_pool2d, complex_sigmoid
+from torch.nn import Module, Parameter, init
+
 from nets.funcs.complexFunctions import complex_dropout, complex_dropout2d, complex_adaptive_avg_pool2d, \
     complex_adaptive_max_pool2d
+from nets.funcs.complexFunctions import complex_relu, complex_max_pool2d, complex_avg_pool2d, complex_sigmoid
 
 
 def apply_complex(fr, fi, input, dtype=torch.complex64):

nets/models/DeepPhys.py

@@ -1,8 +1,8 @@
 import torch
 
 from nets.models.sub_models.AppearanceModel import AppearanceModel_2D
-from nets.models.sub_models.MotionModel import MotionModel
 from nets.models.sub_models.LinearModel import LinearModel
+from nets.models.sub_models.MotionModel import MotionModel
 
 
 class DeepPhys(torch.nn.Module):

nets/models/sub_models/AppearanceModel.py

@@ -1,8 +1,9 @@
 import torch
 from torch.nn import Module
-from nets.modules.modules import DAModule
+
 from nets.blocks.attentionBlocks import AttentionBlock
-from nets.blocks.blocks import EncoderBlock,DecoderBlock
+from nets.blocks.blocks import EncoderBlock, DecoderBlock
+from nets.modules.modules import DAModule
 
 
 class AppearanceModel_2D(Module):

nets/models/sub_models/LinearModel.py

@@ -1,6 +1,7 @@
 import torch
-from nets.layers.complexLayers import ComplexDropout, ComplexLinear
+
 from nets.funcs.complexFunctions import complex_tanh
+from nets.layers.complexLayers import ComplexDropout, ComplexLinear
 
 
 class LinearModel(torch.nn.Module):

nets/models/sub_models/MotionModel.py

@@ -1,5 +1,6 @@
 import torch
-import torch.nn as nn
+
+
 class MotionModel(torch.nn.Module):
     def __init__(self, in_channels, out_channels, kernel_size):
         super().__init__()
@@ -13,11 +14,13 @@ def __init__(self, in_channels, out_channels, kernel_size):
         self.m_dropout1 = torch.nn.Dropout2d(p=0.50)
 
         self.m_avg1 = torch.nn.AvgPool2d(kernel_size=2, stride=2, padding=0)
-        self.m_conv3 = torch.nn.Conv2d(in_channels=out_channels, out_channels=out_channels*2, kernel_size=kernel_size, stride=1,
+        self.m_conv3 = torch.nn.Conv2d(in_channels=out_channels, out_channels=out_channels * 2, kernel_size=kernel_size,
+                                       stride=1,
                                        padding=1)
-        self.m_batch_Normalization3 = torch.nn.BatchNorm2d(out_channels*2)
-        self.m_conv4 = torch.nn.Conv2d(in_channels=out_channels*2, out_channels=out_channels*2, kernel_size=kernel_size, stride=1, padding=1)
-        self.m_batch_Normalization4 = torch.nn.BatchNorm2d(out_channels*2)
+        self.m_batch_Normalization3 = torch.nn.BatchNorm2d(out_channels * 2)
+        self.m_conv4 = torch.nn.Conv2d(in_channels=out_channels * 2, out_channels=out_channels * 2,
+                                       kernel_size=kernel_size, stride=1, padding=1)
+        self.m_batch_Normalization4 = torch.nn.BatchNorm2d(out_channels * 2)
         self.m_dropout2 = torch.nn.Dropout2d(p=0.50)
         self.m_avg2 = torch.nn.AvgPool2d(kernel_size=2, stride=2, padding=0)
 

nets/modules/modules.py

@@ -1,6 +1,8 @@
 import torch
+
 from ..funcs.complexFunctions import complex_matmul
 
+
 class DAModule(torch.nn.Module):
     def __init__(self, in_channels):
         super(DAModule, self).__init__()

optim.py

@@ -1,4 +1,3 @@
-import torch
 import torch.optim as opt
 
 from log import log_warning