Add EER support in Dose Analysis

xmipp3/protocols/protocol_movie_dose_analysis.py

@@ -32,7 +32,7 @@
 
 from pyworkflow import VERSION_3_0
 from pyworkflow.object import Set
-from pyworkflow.protocol.params import (PointerParam, IntParam, FloatParam, LEVEL_ADVANCED)
+from pyworkflow.protocol.params import (PointerParam, IntParam, FloatParam, LEVEL_ADVANCED, GT)
 from pyworkflow.utils.properties import Message
 import pyworkflow.protocol.constants as cons
 from pyworkflow import UPDATED
@@ -42,7 +42,7 @@
 from pwem.objects import SetOfMovies
 from pwem.protocols import ProtProcessMovies
 
-from xmipp3.convert import getScipionObj
+from xmipp3.convert import getScipionObj, isEerMovie
 
 THRESHOLD = 2
 OUTPUT_MOVIES = "outputMovies"
@@ -76,10 +76,15 @@ def __init__(self, **args):
     def _defineParams(self, form):
         form.addSection(label=Message.LABEL_INPUT)
 
+        EER_CONDITION = 'inputMovies is not None and len(inputMovies) > 0 and next(iter(inputMovies.getFiles())).endswith(".eer")'
+
         form.addParam('inputMovies', PointerParam, pointerClass='SetOfMovies',
                       label=Message.LABEL_INPUT_MOVS,
                       help='Select one or several movies. A dose analysis '
                            'be calculated for each one of them.')
+        form.addParam('nFrames', IntParam, label='Number of EER frames',
+                      condition=EER_CONDITION, default=40, validators=[GT(0, "Number of EER frames must be a positive integer (> 0).")],
+                      help='Number of frames to be generated. EER files contain subframes, that will be grouped into the selected number of frames.')
         form.addParam('percentage_threshold', FloatParam, default=5,
                       label="Maximum percentage difference (%)",
                       help='By default, a difference of 5% against the median dose is used to '
@@ -212,12 +217,19 @@ def _processMovies(self, movieIds):
 
 
     def estimatePoissonCount(self, movie):
+
         mean_frames = []
-        n = movie.getNumberOfFrames()
+        filename = movie.getFileName()
+        if filename.endswith(".eer"):
+            n = self.nFrames.get()
+            filename += "#%d,4K,uint16" % n
+        else:
+            n = movie.getNumberOfFrames()
+
         frames = [1, n/2, n]
 
         for frame in frames:
-            frame_image = ImageHandler().read("%d@%s" % (frame, movie.getFileName())).getData()
+            frame_image = ImageHandler().read("%d@%s" % (frame, filename)).getData()
             mean_dose_per_pixel = np.mean(frame_image)
             mean_dose_per_angstrom2 = mean_dose_per_pixel/ self.samplingRate**2
             mean_frames.append(mean_dose_per_angstrom2)