melage 0.0.65__py3-none-any.whl → 1.0.0__py3-none-any.whl

melage/widgets/ApplyMask.py

@@ -1,212 +0,0 @@
-__AUTHOR__ = 'Bahram Jafrasteh'
-
-from PyQt5 import QtCore, QtGui, QtWidgets
-from PyQt5.QtCore import Qt
-from PyQt5.QtCore import pyqtSignal
-
-
-class Worker(QtCore.QObject):
-    """
-
-    """
-    finished = pyqtSignal()
-    def __init__(self, id: int):
-        super().__init__()
-        self.__id = id
-        self.param_set = False
-        self.inputimage = None
-        self.out = None
-    #@QtCore.pyqtSlot()
-    def run(self):
-        if self.inputimage is None or not self.param_set:
-            return
-        from melage.utils.utils import N4_bias_correction
-
-
-        out = N4_bias_correction(self.inputimage, use_otsu=self._outsu,
-                                 shrinkFactor=self._selectedColor, numberFittingLevels=self._filtering_level,
-                                 max_iter=self._max_iter)
-
-
-        self.inputimage = None
-        self.out = out
-        self.param_set = False
-        self.finished.emit()
-        #self.close()
-
-    def set_params(self, inpim, params):
-        self.param_set = True
-        [self._max_iter, self._selectedColor, self._filtering_level, self._outsu] = params
-        self.inputimage = inpim
-
-
-class Masking(QtWidgets.QDialog):
-    """
-    Apply segmentation mask to the image
-    """
-    buttonpressed = pyqtSignal(object)
-    closeSig = pyqtSignal()
-
-    def __init__(self, parent=None):
-        QtWidgets.QDialog.__init__(self, parent)
-        self.setWindowTitle("Child Window!")
-        self.inputimage = None
-
-
-        self.setupUi()
-
-    def setupUi(self):
-        Form = self.window()
-        Form.setObjectName("N4")
-        Form.resize(759, 140)
-        self.grid_main = QtWidgets.QGridLayout(self)
-        self.grid_main.setContentsMargins(0,0,0,0)
-        self.grid_main.setObjectName("gridLayout")
-        self.widget = QtWidgets.QWidget()
-
-        #self.widget.setGeometry(QtCore.QRect(43, 20, 644, 89))
-        self.widget.setObjectName("widget")
-        self.gridLayout = QtWidgets.QGridLayout(self.widget)
-        self.gridLayout.setContentsMargins(10, 10, 10, 10)
-        self.gridLayout.setObjectName("gridLayout")
-
-        self.splitter = QtWidgets.QSplitter(self.widget)
-        self.splitter.setOrientation(QtCore.Qt.Horizontal)
-        self.splitter.setObjectName("splitter")
-
-        #self.gridLayout.addWidget(self.splitter, 0, 1, 1, 1)
-
-        self.mask_color = QtWidgets.QLabel(self.widget)
-        self.mask_color.setAlignment(QtCore.Qt.AlignCenter)
-        self.mask_color.setObjectName("mask_color")
-        self.gridLayout.addWidget(self.mask_color, 0, 2, 1, 1)
-
-        self.mask_keep_remove = QtWidgets.QComboBox(self.widget)
-        self.mask_keep_remove.setObjectName("comboBox_image")
-        self.mask_keep_remove.addItem("Keep")
-        self.mask_keep_remove.addItem("Remove")
-        self.gridLayout.addWidget(self.mask_keep_remove, 0, 1, 1, 1)
-
-        self._combobox_colors = QtWidgets.QComboBox(self.widget)
-        self._combobox_colors.setObjectName("_combobox_colors")
-
-        self.gridLayout.addWidget(self._combobox_colors, 0,3, 1, 3)
-
-
-
-        #self.pushButton_2 = QtWidgets.QPushButton(self.widget)
-        #self.pushButton_2.setObjectName("pushButton")
-        #self.gridLayout.addWidget(self.pushButton_2, 0, 6, 1, 1)
-        self.pushButton = QtWidgets.QPushButton(self.widget)
-        self.pushButton.setObjectName("pushButton_2")
-        self.gridLayout.addWidget(self.pushButton, 1, 5, 1, 3)
-        self.progressBar = QtWidgets.QProgressBar(self.widget)
-        self.progressBar.setEnabled(False)
-        self.progressBar.setProperty("value", 0)
-        self.progressBar.setAlignment(QtCore.Qt.AlignCenter)
-        self.progressBar.setObjectName("progressBar")
-        self.gridLayout.addWidget(self.progressBar, 1, 0, 1, 5)
-        #spacerItem = QtWidgets.QSpacerItem(40, 20, QtWidgets.QSizePolicy.Expanding, QtWidgets.QSizePolicy.Minimum)
-        #self.gridLayout.addItem(spacerItem, 0, 1, 1, 4)
-        self.comboBox_image = QtWidgets.QComboBox(self.widget)
-        self.comboBox_image.setObjectName("comboBox_image")
-        self.comboBox_image.addItem("")
-        self.comboBox_image.addItem("")
-        self.gridLayout.addWidget(self.comboBox_image, 0, 0, 1, 1)
-
-        #self.label_timer = QtWidgets.QLabel(self.widget)
-        #self.label_timer.setAlignment(QtCore.Qt.AlignCenter)
-        #self.label_timer.setObjectName("label")
-        #self.gridLayout.addWidget(self.label_timer, 1, 0, 1, 1)
-        self.grid_main.addWidget(self.widget)
-        self.retranslateUi(Form)
-        QtCore.QMetaObject.connectSlotsByName(Form)
-
-
-
-
-
-        for i in range(4):
-            self._combobox_colors.setItemData(i, Qt.AlignCenter, Qt.TextAlignmentRole)
-            self._combobox_colors.addItem("    fdfdffdff{}    ".format(i+1))
-
-        self.pushButton.clicked.connect(self.accepted_emit)
-
-
-
-
-        self.progressBar.setVisible(True)
-
-
-
-    def setComboBoxColors(self, color_name):
-        self._combobox_colors.clear()
-        for colr in color_name:
-            #self._combobox_colors.setItemData(i, Qt.AlignCenter, Qt.TextAlignmentRole)
-            self._combobox_colors.addItem("   {}   ".format(colr))
-
-    def startTimer(self):
-        self.time_spent = 0
-
-        self.myTimer = QtCore.QTimer(self)
-        self.myTimer.timeout.connect(self.timerTimeout)
-        self.myTimer.start(1000)
-
-    def timerTimeout(self):
-        self.time_spent += 1
-        if self.time_spent == 0:
-            self.time_spent = 50
-        self.label_timer.setText(str(self.time_spent))
-
-
-
-
-    def retranslateUi(self, Form):
-        _translate = QtCore.QCoreApplication.translate
-        Form.setWindowTitle(_translate("Form", "Apply Mask"))
-
-
-
-
-
-
-
-
-
-        self.mask_color.setText(_translate("Form", "Mask Color"))
-        self.mask_keep_remove.setItemText(0, _translate("Form", "        Keep        "))
-        self.mask_keep_remove.setItemText(1, _translate("Form", "        Remove        "))
-        self.comboBox_image.setItemText(0, _translate("Form", "        View 1        "))
-        self.comboBox_image.setItemText(1, _translate("Form", "      View 2      "))
-        #self.label_timer.setText(_translate("Form", "00:00"))
-        self.pushButton.setText(_translate("Form", "Apply"))
-        self.pushButton.setVisible(True)
-    def message(self):
-
-        self.MessageBox.setText('Please fill with numeric (!=0) and text')
-        self.MessageBox.setWindowTitle('Warning')
-        self.MessageBox.show()
-
-    def accepted_emit(self):
-
-        index_image = self.comboBox_image.currentIndex()
-
-        index_keep_remove = self.mask_keep_remove.currentText().strip(' ')
-        keep = False
-        if index_keep_remove.lower()=='keep':
-            keep = True
-        index_color = self._combobox_colors.currentIndex()
-        self.buttonpressed.emit([index_image, index_color, keep])
-
-    def closeEvent(self, a0: QtGui.QCloseEvent) -> None:
-        self.closeSig.emit()
-        super(Masking, self).closeEvent(a0)
-
-def run():
-    import sys
-    app = QtWidgets.QApplication(sys.argv)
-    window = Masking()
-    window.show()
-    sys.exit(app.exec_())
-if __name__ == '__main__':
-    run()

melage/widgets/ChangeSystem.py

@@ -1,152 +0,0 @@
-__AUTHOR__ = 'Bahram Jafrasteh'
-
-
-from PyQt5 import QtCore, QtGui, QtWidgets
-from PyQt5.QtCore import Qt
-from PyQt5.QtCore import pyqtSignal
-
-
-
-
-class ChangeCoordSys(QtWidgets.QDialog):
-    """
-    This class has been implemented to change coordinate system of the current image
-    """
-    buttonpressed = pyqtSignal(object)
-    closeSig = pyqtSignal()
-
-    def __init__(self, parent=None):
-        QtWidgets.QDialog.__init__(self, parent)
-        self.setWindowTitle("Child Window!")
-        self.inputimage = None
-
-
-        self.setupUi()
-
-    def setupUi(self):
-        Form = self.window()
-        Form.setObjectName("N4")
-        #Form.resize(759, 140)
-        self.grid_main = QtWidgets.QGridLayout(self)
-        self.grid_main.setContentsMargins(0,0,0,0)
-        self.grid_main.setObjectName("gridLayout")
-        self.widget = QtWidgets.QWidget()
-        #self.widget.setGeometry(QtCore.QRect(43, 20, 644, 89))
-        self.widget.setObjectName("widget")
-        self.gridLayout = QtWidgets.QGridLayout(self.widget)
-        self.gridLayout.setContentsMargins(10, 10, 10, 10)
-        self.gridLayout.setObjectName("gridLayout")
-
-        self.splitter = QtWidgets.QSplitter(self.widget)
-        self.splitter.setOrientation(QtCore.Qt.Horizontal)
-        self.splitter.setObjectName("splitter")
-
-        #self.gridLayout.addWidget(self.splitter, 0, 1, 1, 1)
-        self._labelfrom = QtWidgets.QLabel(self.widget)
-        self._labelfrom.setAlignment(QtCore.Qt.AlignCenter)
-        self._labelfrom.setObjectName("From")
-        self.gridLayout.addWidget(self._labelfrom, 0, 1, 1, 1)
-
-        self.label_current = QtWidgets.QLabel(self.widget)
-        self.label_current.setText('RAS')
-        self.label_current.setStyleSheet('color: Green')
-        self.gridLayout.addWidget(self.label_current, 0, 2, 1, 1)
-
-        self._labelto = QtWidgets.QLabel(self.widget)
-        self._labelto.setAlignment(QtCore.Qt.AlignCenter)
-        self._labelto.setObjectName("To")
-        self.gridLayout.addWidget(self._labelto, 0, 3, 1, 1)
-        self._combobox_coords = QtWidgets.QComboBox(self.widget)
-        self._combobox_coords.setObjectName("_combobox_coords")
-
-        self.gridLayout.addWidget(self._combobox_coords, 0,4, 1, 2)
-
-
-        #self.pushButton_2 = QtWidgets.QPushButton(self.widget)
-        #self.pushButton_2.setObjectName("pushButton")
-        #self.gridLayout.addWidget(self.pushButton_2, 0, 6, 1, 1)
-        self.pushButton = QtWidgets.QPushButton(self.widget)
-        self.pushButton.setObjectName("pushButton_2")
-        self.gridLayout.addWidget(self.pushButton, 1, 5, 1, 3)
-        self.label_warning = QtWidgets.QLabel(self.widget)
-
-        self.gridLayout.addWidget(self.label_warning, 1, 0, 1, 5)
-        self.label_warning.setStyleSheet('color: Red')
-        self.label_warning.setText('This options is for advanced users')
-
-        #spacerItem = QtWidgets.QSpacerItem(40, 20, QtWidgets.QSizePolicy.Expanding, QtWidgets.QSizePolicy.Minimum)
-        #self.gridLayout.addItem(spacerItem, 0, 1, 1, 4)
-        self.comboBox_image = QtWidgets.QComboBox(self.widget)
-        self.comboBox_image.setObjectName("comboBox_image")
-        self.comboBox_image.addItem("")
-        self.comboBox_image.addItem("")
-        #self.comboBox_image.currentIndexChanged.connect()
-        self.gridLayout.addWidget(self.comboBox_image, 0, 0, 1, 1)
-
-        #self.label_timer = QtWidgets.QLabel(self.widget)
-        #self.label_timer.setAlignment(QtCore.Qt.AlignCenter)
-        #self.label_timer.setObjectName("label")
-        #self.gridLayout.addWidget(self.label_timer, 1, 0, 1, 1)
-
-        self.retranslateUi(Form)
-        QtCore.QMetaObject.connectSlotsByName(Form)
-
-        (('L', 'R'), ('P', 'A'), ('I', 'S'))
-
-        self.systems = ['RAS', 'RSA', 'RPI', 'RPS',
-                   'RIP', 'RIA', 'RSP', 'RSA',
-                    'LAS', 'LAI', 'LPI', 'LPS',
-                   'LSA', 'LIA', 'LIP', 'LSP',
-                   'PIL', 'PIR', 'PSL', 'PSR',
-                   'PLI', 'PRI', 'PLS', 'PRS',
-                   'AIL', 'AIR', 'ASL', 'ASR',
-                   'ALI', 'ARI', 'ALS', 'ARS',
-                   'IPL', 'IPR', 'IAL', 'IAR',
-                   'ILP', 'IRP', 'ILA', 'IRA',
-                   'SPL', 'SPR', 'SAL', 'SAR',
-                   'SLP', 'SRP', 'SLA', 'SRA'
-                   ]
-        for i , ss in enumerate(self.systems):
-            self._combobox_coords.setItemData(i, Qt.AlignCenter, Qt.TextAlignmentRole)
-            self._combobox_coords.addItem(" "*9+ss+" "*9)
-        self.grid_main.addWidget(self.widget)
-        self.pushButton.clicked.connect(self.accepted_emit)
-
-
-
-
-    def retranslateUi(self, Form):
-        _translate = QtCore.QCoreApplication.translate
-        Form.setWindowTitle(_translate("Form", "Change Image Coordinate System"))
-
-
-        self._labelto.setText(_translate("Form", "To"))
-        self._labelfrom.setText(_translate("Form", "From"))
-        self.comboBox_image.setItemText(0, _translate("Form", "        View 1        "))
-        self.comboBox_image.setItemText(1, _translate("Form", "      View 2      "))
-        #self.label_timer.setText(_translate("Form", "00:00"))
-        self.pushButton.setText(_translate("Form", "Apply"))
-        self.pushButton.setVisible(True)
-    def closeEvent(self, a0: QtGui.QCloseEvent) -> None:
-        self.closeSig.emit()
-        super(ChangeCoordSys, self).closeEvent(a0)
-
-    def accepted_emit(self):
-
-
-        index_image = self.comboBox_image.currentIndex()
-
-        index_color = self._combobox_coords.currentIndex()
-        CS = self.systems[index_color]
-        self.buttonpressed.emit([index_image, CS])
-
-
-
-def run():
-    import sys
-    app = QtWidgets.QApplication(sys.argv)
-    window = ChangeCoordSys()
-    window.show()
-    sys.exit(app.exec_())
-if __name__ == '__main__':
-    run()

melage/widgets/DeepLModels/NPP/dataset/mri_dataset_affine.py

@@ -1,149 +0,0 @@
-import os.path
-from torch.utils.data import Dataset
-import nibabel as nib
-from multiprocessing import Manager
-import random
-import torchio as tio
-import pickle
-import torch
-
-error_list = [5298, 5894,  393,   66,   68, 6576,   76, 6653,  347, 6218,  804,  844,
-        5445,  751, 6075, 5501,  368,  236,  269, 5472,  494, 6711, 5571, 6543,
-        5837, 5586, 6663,  791, 6113,  318,  299,  688,  323,   26, 6496, 5746]
-def Generate_dataset():
-
-    # This if else sentence is used to decide whether enables dataset cache funciton
-    if False:
-        cache = DatasetCache(None,use_cache=False)
-        manager2 = Manager()
-        cache2 = DatasetCache(manager2,use_cache=True)
-    else:
-        manager = None
-        cache = DatasetCache(manager,use_cache=False)
-        manager2 = None
-        cache2 = DatasetCache(manager2,use_cache=False)
-
-    # Loading all availible files
-    if os.path.exists('dataset/pretrain_files_list.pkl'):
-        with open('dataset/pretrain_files_list.pkl', 'rb') as file:
-            files_list = pickle.load(file)
-
-        # splitting the training and testing datasets. The Oasis will be treated as testing while the rest of them will be trated as training dataset.
-        testing_set = files_list[0:1]
-
-        training_sets = files_list
-        del training_sets[0]
-
-        training_sets[2] = [i.replace('GSP','GSP/FS_4p5') for i in training_sets[2]]
-        training_sets_path = sum(training_sets, [])
-        for index in sorted(error_list, reverse=True):
-            del training_sets_path[index]
-        training_sets_path= ['/scratch/datasets/xh278/orig'+i for i in training_sets_path]
-        testing_set_path = sum(testing_set, [])
-
-        train_dataset = VoxelDataset(training_sets_path, cache=cache, train=True)
-        val_dataset = VoxelDataset(testing_set_path, cache=cache2, train=False)
-
-    return train_dataset,val_dataset
-
-class DatasetCache(object):
-    def __init__(self, manager, use_cache=True):
-        self.use_cache = use_cache
-        self.manager = manager
-        if self.manager is not None:
-            self._dict = manager.dict()
-
-    def is_cached(self, key):
-        if not self.use_cache:
-            return False
-        return str(key) in self._dict
-
-    def reset(self):
-        self._dict.clear()
-
-    def get(self, key):
-        if not self.use_cache:
-            raise AttributeError('Data caching is disabled and get funciton is unavailable! Check your config.')
-        return self._dict[str(key)]
-
-    def cache(self, key, img, lbl):
-        # only store if full data in memory is enabled
-        if not self.use_cache:
-            return
-        # only store if not already cached
-        if str(key) in self._dict:
-            return
-        self._dict[str(key)] = (img, lbl)
-
-
-class VoxelDataset(Dataset):
-    def __init__(self, norm_file_path,  rescale_sdf=True, cache = None,train=True):
-        # note that input image paths are for processed images rather than unprocessed
-        self.norm_file_path = [i.replace('proc/ForAdrian_Talairach','orig').replace('talairach/','') for i in norm_file_path]
-        self.orig_file_path = [i.replace('norm.mgz','orig.mgz').replace('talairach/','') for i in self.norm_file_path]
-
-        self.intensity_spatial_norm_file_path = [i.replace('orig.mgz','talairach/norm.mgz') for i in self.orig_file_path]
-        self.intensity_norm_file_path =[i.replace('orig.mgz','norm.mgz') for i in self.orig_file_path]
-        self.seg_path =[i.replace('orig.mgz','talairach/aseg.mgz') for i in self.orig_file_path]
-
-        removal = []
-        for i in range(len(self.seg_path)):
-            if not os.path.exists(self.seg_path[i]):
-                removal.append(i)
-        removal = []
-
-        for i in reversed(removal):
-            self.seg_path.pop(i)
-            self.intensity_spatial_norm_file_path.pop(i)
-            self.orig_file_path.pop(i)
-        self.rescale_sdf = rescale_sdf
-        self.cache = cache
-        self.train = train
-        self.transform = tio.Compose([
-            tio.RandomGamma(log_gamma = (-0.3,0.3),p=0.3),
-            ]
-        )
-    def __len__(self):
-        return len(self.orig_file_path)
-
-    def __getitem__(self, index):
-        image_resolution= 256
-        normalization=255
-        output_index = [index]
-        for cur_index in output_index:
-            array = nib.load(self.orig_file_path[cur_index])
-            input = array.get_fdata()
-            input = torch.FloatTensor(input)
-            input = input.unsqueeze(0).unsqueeze(0)
-            if image_resolution != 256:
-                input = torch.nn.functional.interpolate(input, size=[image_resolution, image_resolution, image_resolution], mode='trilinear',align_corners=False)[0]
-            else:
-                input = torch.Tensor(input)[0]
-
-            array = nib.load(self.intensity_spatial_norm_file_path[cur_index])
-            intensity_spatial_norm = array.get_fdata()
-            intensity_spatial_norm = torch.FloatTensor(intensity_spatial_norm)
-            intensity_spatial_norm = intensity_spatial_norm.unsqueeze(0).unsqueeze(0)
-            if image_resolution != 256:
-                intensity_spatial_norm = torch.nn.functional.interpolate(intensity_spatial_norm, size=[image_resolution, image_resolution, image_resolution], mode='trilinear',align_corners=False)[0]
-            else:
-                intensity_spatial_norm = torch.Tensor(intensity_spatial_norm)[0]
-
-            seg = intensity_spatial_norm.unsqueeze(0)
-            if image_resolution != 256:
-                seg = torch.nn.functional.interpolate(seg, size=[image_resolution, image_resolution, image_resolution], mode='nearest')[0]
-            else:
-                seg = torch.Tensor(seg)[0]
-        if self.train:
-            subject = tio.Subject(image = tio.ScalarImage(tensor=input))
-            transformed = self.transform(subject)
-            input = transformed['image'].data
-            input = input.clip(0, normalization) / normalization
-            intensity_spatial_norm = intensity_spatial_norm.clip(0, normalization) / normalization
-        else:
-            input = input.clip(0, normalization) / normalization
-            intensity_spatial_norm = intensity_spatial_norm.clip(0, normalization) / normalization
-        return input,intensity_spatial_norm,seg,cur_index
-
-
-

melage/widgets/DeepLModels/NPP/models/losses.py

@@ -1,146 +0,0 @@
-import torch
-import torch.nn.functional as F
-from torch.autograd import Variable
-import numpy as np
-from math import exp
-
-def gaussian(window_size, sigma):
-    gauss = torch.Tensor([exp(-(x - window_size//2)**2/float(2*sigma**2)) for x in range(window_size)])
-    return gauss/gauss.sum()
-
-def create_window(window_size, channel):
-    _1D_window = gaussian(window_size, 1.5).unsqueeze(1)
-    _2D_window = _1D_window.mm(_1D_window.t()).float().unsqueeze(0).unsqueeze(0)
-    window = Variable(_2D_window.expand(channel, 1, window_size, window_size).contiguous())
-    return window
-
-def create_window_3D(window_size, channel):
-    _1D_window = gaussian(window_size, 1.5).unsqueeze(1)
-    _2D_window = _1D_window.mm(_1D_window.t())
-    _3D_window = _1D_window.mm(_2D_window.reshape(1, -1)).reshape(window_size, window_size, window_size).float().unsqueeze(0).unsqueeze(0)
-    window = Variable(_3D_window.expand(channel, 1, window_size, window_size, window_size).contiguous())
-    return window
-
-def _ssim(img1, img2, window, window_size, channel, size_average = True):
-    mu1 = F.conv2d(img1, window, padding = window_size//2, groups = channel)
-    mu2 = F.conv2d(img2, window, padding = window_size//2, groups = channel)
-
-    mu1_sq = mu1.pow(2)
-    mu2_sq = mu2.pow(2)
-    mu1_mu2 = mu1*mu2
-
-    sigma1_sq = F.conv2d(img1*img1, window, padding = window_size//2, groups = channel) - mu1_sq
-    sigma2_sq = F.conv2d(img2*img2, window, padding = window_size//2, groups = channel) - mu2_sq
-    sigma12 = F.conv2d(img1*img2, window, padding = window_size//2, groups = channel) - mu1_mu2
-
-    C1 = 0.01**2
-    C2 = 0.03**2
-
-    ssim_map = ((2*mu1_mu2 + C1)*(2*sigma12 + C2))/((mu1_sq + mu2_sq + C1)*(sigma1_sq + sigma2_sq + C2))
-
-    if size_average:
-        return ssim_map.mean()
-    else:
-        return ssim_map.mean(1).mean(1).mean(1)
-
-def f_ssim_3D(img1, img2, window, window_size, channel, size_average = True):
-    #window = torch.ones(window.shape).type_as(img1)
-    #window = window/window.sum()
-    mu1 = F.conv3d(img1, window, padding = window_size//2, groups = channel)
-    mu2 = F.conv3d(img2, window, padding = window_size//2, groups = channel)
-
-    mu1_sq = mu1.pow(2)
-    mu2_sq = mu2.pow(2)
-
-    mu1_mu2 = mu1*mu2
-
-    sigma1_sq = F.conv3d(img1*img1, window, padding = window_size//2, groups = channel) - mu1_sq
-    sigma2_sq = F.conv3d(img2*img2, window, padding = window_size//2, groups = channel) - mu2_sq
-    sigma12 = F.conv3d(img1*img2, window, padding = window_size//2, groups = channel) - mu1_mu2
-
-    C1 = 0.01**2
-    C2 = 0.03**2
-
-    ssim_map = ((2*mu1_mu2 + C1)*(2*sigma12 + C2))/((mu1_sq + mu2_sq + C1)*(sigma1_sq + sigma2_sq + C2))
-
-    if size_average:
-        return ssim_map.mean()
-    else:
-        return ssim_map.mean(1).mean(1).mean(1)
-
-
-
-class SSIM(torch.nn.Module):
-    def __init__(self, window_size = 11, size_average = True):
-        super(SSIM, self).__init__()
-        self.window_size = window_size
-        self.size_average = size_average
-        self.channel = 1
-        self.window = create_window(window_size, self.channel)
-
-    def forward(self, img1, img2):
-        (_, channel, _, _) = img1.size()
-
-        if channel == self.channel and self.window.data.type() == img1.data.type():
-            window = self.window
-        else:
-            window = create_window(self.window_size, channel)
-
-            if img1.is_cuda:
-                window = window.cuda(img1.get_device())
-            window = window.type_as(img1)
-
-            self.window = window
-            self.channel = channel
-
-
-        return _ssim(img1, img2, window, self.window_size, channel, self.size_average)
-
-
-class SSIM3D(torch.nn.Module):
-    def __init__(self, window_size = 11, size_average = True):
-        super(SSIM3D, self).__init__()
-        self.window_size = window_size
-        self.size_average = size_average
-        self.channel = 1
-        self.window = create_window_3D(window_size, self.channel)
-
-    def forward(self, img1, img2):
-        (_, channel, _, _, _) = img1.size()
-
-        if channel == self.channel and self.window.data.type() == img1.data.type():
-            window = self.window
-        else:
-            window = create_window_3D(self.window_size, channel)
-
-            if img1.is_cuda:
-                window = window.cuda(img1.get_device())
-            window = window.type_as(img1)
-
-            self.window = window
-            self.channel = channel
-
-
-        return f_ssim_3D(img1, img2, window, self.window_size, channel, self.size_average)
-
-
-def ssim(img1, img2, window_size = 11, size_average = True):
-    (_, channel, _, _) = img1.size()
-    window = create_window(window_size, channel)
-
-    if img1.is_cuda:
-        window = window.cuda(img1.get_device())
-    window = window.type_as(img1)
-
-    return _ssim(img1, img2, window, window_size, channel, size_average)
-
-def ssim3D(img1, img2, window_size = 11, size_average = True):
-    (_, channel, _, _, _) = img1.size()
-    window = create_window_3D(window_size, channel)
-
-    if img1.is_cuda:
-        window = window.cuda(img1.get_device())
-    window = window.type_as(img1)
-
-    return f_ssim_3D(img1, img2, window, window_size, channel, size_average)
-