Jj update to latest version

CHANGES.txt

@@ -0,0 +1,2 @@
+3.1.0:
+    - Tensorflow 2 version. Source code version 0.2.1.

cryocare/__init__.py

@@ -32,11 +32,10 @@
 
 _logo = "icon.png"
 _references = ['buchholz2019cryo', 'buchholz2019content']
-__version__ = "3.0.1"
+__version__ = "3.1.0"
 
 
 class Plugin(pwem.Plugin):
-
     _homeVar = CRYOCARE_HOME
     _url = 'https://github.com/scipion-em/scipion-em-cryocare'
 
@@ -72,21 +71,16 @@ def defineBinaries(cls, env):
 
         # Create the environment
         installationCmd += 'conda create -y -n %s -c conda-forge -c anaconda python=3.8 ' \
-                           'cudnn=7.6.5=cuda10.1_0 && ' % CRYOCARE_ENV_NAME
+                           'cudatoolkit=11.0 cudnn=8.0 && ' % CRYOCARE_ENV_NAME
+        # 'keras-gpu=2.3.1 ' \
 
         # Activate new the environment
         installationCmd += 'conda activate %s && ' % CRYOCARE_ENV_NAME
 
-        # Install non-conda required packages
-        installationCmd += 'pip install tensorflow-gpu==2.3.3 && '
-        installationCmd += 'pip install mrcfile && '
-        installationCmd += 'pip install csbdeep '
-        # I had the same issue and was able to fix this by setting h5py < 3.0.0.
-        # Looks like here was a 3.0 release of h5py recently where they changed how strings are stored/read.
-        # https://github.com/keras-team/keras/issues/14265
-
-        # Install cryoCARE
-        installationCmd += 'pip install %s==%s &&' % (CRYOCARE, CRYOCARE_DEFAULT_VERSION)
+        # Install cryoCARE and the rest of dependencies
+        installationCmd += 'virtualenv==20.3.1 && '
+        installationCmd += 'pip install tensorflow-gpu==2.4.0 && '
+        installationCmd += 'pip install %s==%s && ' % (CRYOCARE, CRYOCARE_DEFAULT_VERSION)
 
         # Flag installation finished
         installationCmd += 'touch %s' % CRYOCARE_INSTALLED

cryocare/constants.py

@@ -26,8 +26,8 @@
 # **************************************************************************
 
 CRYOCARE_HOME = 'CRYOCARE_HOME'
-V0_1_1 = '0.1.1'
-CRYOCARE_DEFAULT_VERSION = V0_1_1
+V0_2_1 = '0.2.1'
+CRYOCARE_DEFAULT_VERSION = V0_2_1
 CRYOCARE = 'cryoCARE'
 CRYOCARE_ENV_NAME = '%s-%s' % (CRYOCARE, CRYOCARE_DEFAULT_VERSION)
 CRYOCARE_ENV_ACTIVATION = 'CRYOCARE_ENV_ACTIVATION'
@@ -40,4 +40,5 @@
 MEAN_STD_FN = 'mean_std.npz'
 TRAIN_DATA_CONFIG = 'training_data_config'
 CRYOCARE_MODEL = 'cryoCARE_model'
+CRYOCARE_MODEL_TGZ = CRYOCARE_MODEL + '.tar.gz'
 PREDICT_CONFIG = 'predict_config'

cryocare/objects.py

@@ -1,5 +1,7 @@
+from os.path import join
 
 import pyworkflow.object as pwobj
+from cryocare.constants import CRYOCARE_MODEL
 from pwem import EMObject
 
 
@@ -20,13 +22,13 @@ def __str__(self):
 
 
 class CryocareModel(EMObject):
-    def __init__(self, basedir=None, train_data_dir=None, **kwargs):
+    def __init__(self, model_file=None, train_data_dir=None, **kwargs):
         EMObject.__init__(self, **kwargs)
-        self._basedir = pwobj.String(basedir)
+        self._model_file = pwobj.String(model_file)
         self._train_data_dir = pwobj.String(train_data_dir)
 
     def getPath(self):
-        return self._basedir.get()
+        return self._model_file.get()
 
     def getTrainDataDir(self):
         return self._train_data_dir.get()

cryocare/protocols/protocol_load_model.py

@@ -1,21 +1,26 @@
-import glob
+from enum import Enum
 from os.path import exists, join
 
+from cryocare.utils import makeDatasetSymLinks, getModelName
 from pwem.protocols import EMProtocol
 from pyworkflow import BETA
 from pyworkflow.protocol import PathParam, FileParam
 from pyworkflow.utils import Message, createLink
 
-from cryocare.constants import TRAIN_DATA_FN, VALIDATION_DATA_FN, CRYOCARE_MODEL
+from cryocare.constants import TRAIN_DATA_FN, VALIDATION_DATA_FN
 from cryocare.objects import CryocareModel
-from cryocare.utils import makeDatasetSymLinks
+
+
+class outputObjects(Enum):
+    model = CryocareModel
 
 
 class ProtCryoCARELoadModel(EMProtocol):
     """Use two data-independent reconstructed tomograms to train a 3D cryo-CARE network."""
 
     _label = 'CryoCARE Load Model'
     _devStatus = BETA
+    _possibleOutputs = outputObjects
 
     # -------------------------- DEFINE param functions ----------------------
     def _defineParams(self, form):
@@ -25,11 +30,16 @@ def _defineParams(self, form):
         """
         # You need a params to belong to a section:
         form.addSection(label=Message.LABEL_INPUT)
-        form.addParam('basedir', PathParam,
-                      label='Base directory of the trained cryoCARE model',
+        form.addParam('trainDataModel', PathParam,
+                      label='Pre-trained cryoCARE model (.tar.gz)',
                       important=True,
                       allowsNull=False,
-                      help='It must contain a model in .h5 format.')
+                      help='It is a .tar.gz file containing a folder that contains, in turn, the following files:\n\n'
+                           '\t- config.json\n'
+                           '\t- history.dat\n'
+                           '\t- norm.json\n'
+                           '\t- weights_best.h5\n'
+                           '\t- weights_last.h5\n')
         form.addParam('trainDataDir', FileParam,
                       label='Directory of the prepared data for training',
                       important=True,
@@ -46,19 +56,17 @@ def _initialize(self):
         # model, but they are located in the training data generation extra directory. Hence, a symbolic link will
         # be created
         makeDatasetSymLinks(self, self.trainDataDir.get())
-        createLink(join('..', self.basedir.get()), self._getExtraPath(CRYOCARE_MODEL))
+        createLink(join(self.trainDataModel.get()), getModelName(self))
 
     def createOutputStep(self):
-        model = CryocareModel(basedir=self._getExtraPath(), train_data_dir=self._getExtraPath())
-        self._defineOutputs(model=model)
+        model = CryocareModel(model_file=getModelName(self), train_data_dir=self._getExtraPath())
+        self._defineOutputs(**{outputObjects.model.name: model})
 
     # --------------------------- INFO functions -----------------------------------
     def _validate(self):
         errors = []
-        if not exists(self.basedir.get()):
-            errors.append('Training model base directory does not exists.')
-        elif not glob.glob(join(self.basedir.get(), '*.h5')):
-            errors.append('No model files were found in the introduced training model base directory.')
+        if not exists(self.trainDataModel.get()):
+            errors.append('Training model introduced does not exists.')
 
         if not exists(self.trainDataDir.get()):
             errors.append('Directory of the prepared data for training does not exists.')
@@ -75,5 +83,6 @@ def _summary(self):
         summary = []
 
         if self.isFinished():
-            summary.append("Loaded training model_dir = *%s*" % self.basedir.get())
+            summary.append("Loaded training model_dir = *%s*" % self.trainDataModel.get())
         return summary
+

cryocare/protocols/protocol_predict.py

@@ -1,28 +1,65 @@
+# **************************************************************************
+# *
+# * Authors:     Scipion Team
+# *
+# * Unidad de  Bioinformatica of Centro Nacional de Biotecnologia , CSIC
+# *
+# * This program is free software; you can redistribute it and/or modify
+# * it under the terms of the GNU General Public License as published by
+# * the Free Software Foundation; either version 2 of the License, or
+# * (at your option) any later version.
+# *
+# * This program is distributed in the hope that it will be useful,
+# * but WITHOUT ANY WARRANTY; without even the implied warranty of
+# * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# * GNU General Public License for more details.
+# *
+# * You should have received a copy of the GNU General Public License
+# * along with this program; if not, write to the Free Software
+# * Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA
+# * 02111-1307  USA
+# *
+# *  All comments concerning this program package may be sent to the
+# *  e-mail address '[email protected]'
+# *
+# **************************************************************************
+import glob
 import json
-from os.path import abspath, join
-
+import re
+import shutil
+from enum import Enum
+from os.path import join
+from cryocare.utils import checkInputTomoSetsSize
 from pwem.protocols import EMProtocol
 from pyworkflow import BETA
 from pyworkflow.protocol import params, StringParam
-from pyworkflow.utils import Message, removeBaseExt, makePath
+from pyworkflow.utils import Message, makePath
 from scipion.constants import PYTHON
-
 from cryocare import Plugin
-from tomo.objects import Tomogram
-from tomo.protocols import ProtTomoBase
+from tomo.objects import Tomogram, SetOfTomograms
+from cryocare.constants import PREDICT_CONFIG
+
 
-from cryocare.constants import PREDICT_CONFIG, CRYOCARE_MODEL
-from cryocare.utils import CryocareUtils as ccutils
+DENOISED_SUFFIX = 'denoised'
+EVEN = 'even'
 
 
-class ProtCryoCAREPrediction(EMProtocol, ProtTomoBase):
+class outputObjects(Enum):
+    tomograms = SetOfTomograms
+
+
+class ProtCryoCAREPrediction(EMProtocol):
     """Generate the final restored tomogram by applying the cryoCARE trained network to both
 tomograms followed by per-pixel averaging."""
 
     _label = 'CryoCARE Prediction'
-    _configPath = []
-    _outputFiles = []
     _devStatus = BETA
+    _possibleOutputs = outputObjects
+
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+        self._configPath = {}
+        self.sRate = None
 
     # -------------------------- DEFINE param functions ----------------------
     def _defineParams(self, form):
@@ -70,70 +107,94 @@ def _defineParams(self, form):
 
     # --------------------------- STEPS functions ------------------------------
     def _insertAllSteps(self):
-        numTomo = 0
-        makePath(self._getPredictConfDir())
+        self._initialize()
         # Insert processing steps
         for evenTomo, oddTomo in zip(self.even.get(), self.odd.get()):
-            self._insertFunctionStep(self.preparePredictStep, evenTomo.getFileName(), oddTomo.getFileName(), numTomo)
-            self._insertFunctionStep(self.predictStep, numTomo)
-            numTomo += 1
+            tsId = evenTomo.getTsId()
+            self._insertFunctionStep(self.preparePredictStep, evenTomo.getFileName(), oddTomo.getFileName(), tsId)
+            self._insertFunctionStep(self.predictStep, tsId)
+            self._insertFunctionStep(self.createOutputStep, tsId)
 
-        self._insertFunctionStep(self.createOutputStep)
+    def _initialize(self):
+        makePath(self._getPredictConfDir())
+        self.sRate = self.even.get().getSamplingRate()
 
-    def preparePredictStep(self, evenTomo, oddTomo, numTomo):
-        outputName = self._getOutputName(evenTomo)
-        self._outputFiles.append(outputName)
+    def preparePredictStep(self, evenTomo, oddTomo, tsId):
         config = {
-            'model_name': CRYOCARE_MODEL,
             'path': self.model.get().getPath(),
             'even': evenTomo,
             'odd': oddTomo,
-            'output_name': outputName,
-            'n_tiles': [int(i) for i in self.n_tiles.get().split()]
+            'n_tiles': [int(i) for i in self.n_tiles.get().split()],
+            'output': self._getOutputPath(tsId),
+            'overwrite': False
         }
-        self._configPath.append(join(self._getPredictConfDir(), '{}_{:03d}.json'.format(PREDICT_CONFIG, numTomo)))
-        with open(self._configPath[numTomo], 'w+') as f:
+        self._configPath[tsId] = join(self._getPredictConfDir(), '%s_%s.json' % (PREDICT_CONFIG, tsId))
+        with open(self._configPath[tsId], 'w+') as f:
             json.dump(config, f, indent=2)
 
-    def predictStep(self, numTomo):
+    def predictStep(self, tsId):
         # Run cryoCARE
-        Plugin.runCryocare(self, PYTHON, '$(which cryoCARE_predict.py) --conf %s' % self._configPath[numTomo],
+        Plugin.runCryocare(self, PYTHON, '$(which cryoCARE_predict.py) --conf %s' % self._configPath[tsId],
                            gpuId=getattr(self, params.GPU_LIST).get())
+        # Remove even/odd words from the output name to avoid confusion
+        origName = self._getOutputFile(tsId)
+        finalNameRe = re.compile(re.escape(EVEN), re.IGNORECASE)  # Used to do a case-insensitive replacement
+        shutil.move(origName, finalNameRe.sub('', origName))
 
-    def createOutputStep(self):
-        outputSetOfTomo = self._createSetOfTomograms(suffix='_denoised')
-        outputSetOfTomo.copyInfo(self.even.get())
+    def createOutputStep(self, tsId):
+        outputSetOfTomo = getattr(self, outputObjects.tomograms.name, None)
+        if not outputSetOfTomo:
+            outputSetOfTomo = SetOfTomograms.create(self._getPath(), template='tomograms%s.sqlite', suffix=DENOISED_SUFFIX)
+            outputSetOfTomo.copyInfo(self.even.get())
 
-        for i, inTomo in enumerate(self.even.get()):
-            tomo = Tomogram()
-            tomo.setLocation(self._outputFiles[i])
-            tomo.setSamplingRate(inTomo.getSamplingRate())
-            outputSetOfTomo.append(tomo)
+        tomo = self._genOutputTomogram(tsId)
+        outputSetOfTomo.append(tomo)
 
-        self._defineOutputs(outputTomograms=outputSetOfTomo)
+        self._defineOutputs(**{outputObjects.tomograms.name: outputSetOfTomo})
+        self._defineSourceRelation(self.even.get(), outputSetOfTomo)
+        self._defineSourceRelation(self.odd.get(), outputSetOfTomo)
+        self._defineSourceRelation(self.model.get(), outputSetOfTomo)
 
     # --------------------------- INFO functions -----------------------------------
     def _summary(self):
         """ Summarize what the protocol has done"""
         summary = []
 
         if self.isFinished():
-            summary.append(
-                "Tomogram denoising finished.")
+            summary.append("Tomogram denoising finished.")
         return summary
 
     def _validate(self):
         validateMsgs = []
-
-        msg = ccutils.checkInputTomoSetsSize(self.even.get(), self.odd.get())
+        # Check the sampling rate
+        sRateEven = self.even.get().getSamplingRate()
+        sRateOdd = self.odd.get().getSamplingRate()
+        if sRateEven != sRateOdd:
+            validateMsgs.append('The sampling rate of the introduced sets of tomograms is different:\n'
+                                'Even SR %.2f != Odd SR %.2f\n\n' % (sRateEven, sRateOdd))
+        # Check the size
+        msg = checkInputTomoSetsSize(self.even.get(), self.odd.get())
         if msg:
-            validateMsgs.append()
+            validateMsgs.append(msg)
         return validateMsgs
 
     # --------------------------- UTIL functions -----------------------------------
-    def _getOutputName(self, inTomoName):
-        outputName = removeBaseExt(inTomoName) + '_denoised.mrc'
-        return abspath(self._getExtraPath(outputName.replace('_Even', '').replace('_Odd', '')))
-
     def _getPredictConfDir(self):
         return self._getExtraPath(PREDICT_CONFIG)
+
+    def _getOutputPath(self, tsId):
+        """cryoCARE will generate a new folder for each tomogram denoised. Apart from that, if the
+        tomograms were imported, the 'Even_' word can be included in the tsId, as in that case it will be
+        the filename. To avoid confusion, it's removed from the generated folder name."""
+        outPath = self._getExtraPath(tsId + '_' + DENOISED_SUFFIX)
+        outPathRe = re.compile(re.escape(EVEN), re.IGNORECASE)  # Used to carry out a case-insensitive replacement
+        return outPathRe.sub('', outPath)
+
+    def _getOutputFile(self, tsId):
+        return glob.glob(join(self._getOutputPath(tsId), '*.mrc'))[0] # Only one file is contained in each dir
+
+    def _genOutputTomogram(self, tsId):
+        tomo = Tomogram()
+        tomo.setLocation(self._getOutputFile(tsId))
+        tomo.setSamplingRate(self.sRate)
+        return tomo