# **************************************************************************
# *
# * Authors: J.M. De la Rosa Trevin (jmdelarosa@cnb.csic.es)
# * Vahid Abrishami (vabrishami@cnb.csic.es)
# * Josue Gomez Blanco (josue.gomez-blanco@mcgill.ca)
# * David Strelak (davidstrelak@gmail.com)
# *
# * 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 'scipion@cnb.csic.es'
# *
# **************************************************************************
import os.path
import numpy as np
from math import ceil
import pyworkflow.utils as pwutils
from pyworkflow.utils import yellowStr
import pyworkflow.object as pwobj
import pyworkflow.protocol.params as params
import pyworkflow.protocol.constants as cons
import pwem.emlib.metadata as md
from pwem import emlib
from pwem.objects import Image
from pwem.protocols.protocol_align_movies import createAlignmentPlot
from pyworkflow import VERSION_1_1
from pwem.protocols import ProtAlignMovies
from xmipp3.convert import writeMovieMd
from xmipp3.base import isXmippCudaPresent
import xmipp3.utils as xmutils
[docs]class XmippProtMovieCorr(ProtAlignMovies):
"""
Wrapper protocol to Xmipp Movie Alignment by cross-correlation
"""
OUTSIDE_WRAP = 0
OUTSIDE_AVG = 1
OUTSIDE_VALUE = 2
INTERP_LINEAR = 0
INTERP_CUBIC = 1
# Map to xmipp interpolation values in command line
INTERP_MAP = {INTERP_LINEAR: 1, INTERP_CUBIC: 3}
NO_ROTATION = 0
NO_FLIP = 0
_label = 'FlexAlign'
_lastUpdateVersion = VERSION_1_1
def __init__(self, **args):
ProtAlignMovies.__init__(self, **args)
self.stepsExecutionMode = cons.STEPS_PARALLEL
#--------------------------- DEFINE param functions ------------------------
def _defineAlignmentParams(self, form):
ProtAlignMovies._defineAlignmentParams(self, form)
form.addParam('splineOrder', params.EnumParam, condition="doSaveAveMic or doSaveMovie",
default=self.INTERP_CUBIC, choices=['linear', 'cubic'],
expertLevel=cons.LEVEL_ADVANCED,
label='Interpolation',
help="linear (faster but lower quality), "
"cubic (slower but more accurate).")
form.addHidden(params.USE_GPU, params.BooleanParam, default=True,
label="Use GPU for execution",
help="This protocol has both CPU and GPU implementation.\
Select the one you want to use.")
form.addHidden(params.GPU_LIST, params.StringParam, default='0',
expertLevel=cons.LEVEL_ADVANCED,
label="Choose GPU IDs",
help="Add a list of GPU devices that can be used")
form.addParam('maxResForCorrelation', params.FloatParam, default=30,
label='Maximal resolution (A)',
help="Maximal resolution in A that will be preserved during correlation.")
form.addParam('doComputePSD', params.BooleanParam, default=True,
label="Compute PSD (before/after)?",
help="If Yes, the protocol will compute for each movie "
"the PSD of the average micrograph (without CC "
"alignement) and after that, to compare each PSDs")
form.addParam('maxShift', params.IntParam, default=30,
expertLevel=cons.LEVEL_ADVANCED,
label="Maximum shift (pixels)",
help='Maximum allowed distance (in pixels) that each '
'frame can be shifted with respect to the next.')
form.addParam('outsideMode', params.EnumParam,
choices=['Wrapping','Average','Value'],
default=self.OUTSIDE_WRAP,
expertLevel=cons.LEVEL_ADVANCED,
label="How to fill borders",
help='How to fill the borders when shifting the frames')
# this must stay together with the outside mode
form.addParam('outsideValue', params.FloatParam, default=0.0,
expertLevel=cons.LEVEL_ADVANCED,
condition="outsideMode==2",
label='Fill value',
help="Fixed value for filling borders")
#Local alignment params
group = form.addGroup('Local alignment')
group.addParam('doLocalAlignment', params.BooleanParam, default=True,
label="Compute local alignment?",
help="If Yes, the protocol will try to determine local shifts, similarly to MotionCor2.")
group.addParam('autoControlPoints', params.BooleanParam, default=True,
label="Auto control points",
expertLevel=cons.LEVEL_ADVANCED,
condition='doLocalAlignment',
help="If on, protocol will automatically determine necessary number of control points.")
line = group.addLine('Number of control points',
expertLevel=cons.LEVEL_ADVANCED,
help='Number of control points use for BSpline.',
condition='not autoControlPoints')
line.addParam('controlPointX', params.IntParam, default=6, label='X')
line.addParam('controlPointY', params.IntParam, default=6, label='Y')
line.addParam('controlPointT', params.IntParam, default=5, label='t')
group.addParam('minLocalRes', params.FloatParam, default=500,
expertLevel=cons.LEVEL_ADVANCED,
label='Min size of the patch (A)',
help="How many A should contain each patch?")
group.addParam('skipAutotuning', params.BooleanParam, default=False,
expertLevel=cons.LEVEL_ADVANCED,
label='Skip autotuning',
help="We try to faster settings of for the FFT library. This takes some time, "
"but consecutive executions will be faster and use less memory."
"Set to True (autotuning will be turned off) if you process just few movies")
group.addParam('groupNFrames', params.IntParam, default=3,
expertLevel=cons.LEVEL_ADVANCED,
label='Group N frames',
help='Group every specified number of frames by adding them together. \
The alignment is then performed on the summed frames.',
condition='doLocalAlignment')
form.addSection(label="Gain orientation")
form.addParam('gainRot', params.EnumParam,
choices=['no rotation', '90 degrees',
'180 degrees', '270 degrees'],
label="Rotate gain reference:",
default=self.NO_ROTATION,
display=params.EnumParam.DISPLAY_COMBO,
help="Rotate gain reference counter-clockwise.")
form.addParam('gainFlip', params.EnumParam,
choices=['no flip', 'upside down', 'left right'],
label="Flip gain reference:", default=self.NO_FLIP,
display=params.EnumParam.DISPLAY_COMBO,
help="Flip gain reference after rotation. "
"For tiff movies, gain is automatically upside-down flipped")
form.addParallelSection(threads=1, mpi=1)
#--------------------------- STEPS functions -------------------------------
def _processMovie(self, movie):
try:
self.tryProcessMovie(movie)
except Exception as ex:
print(yellowStr("We cannot process %s" % movie.getFileName()))
print(ex)
[docs] def getOutsideModeArg(self):
if self.outsideMode == self.OUTSIDE_WRAP:
return ' --outside wrap'
if self.outsideMode == self.OUTSIDE_AVG:
return ' --outside avg'
if self.outsideMode == self.OUTSIDE_VALUE:
return ' --outside value %f' % self.outsideValue
return ''
[docs] def getCropCornerArg(self, x, y):
# Assume that if you provide one cropDim, you provide all
offsetX = self.cropOffsetX.get()
offsetY = self.cropOffsetY.get()
cropDimX = self.cropDimX.get()
cropDimY = self.cropDimY.get()
if 0 == offsetX and 0 == offsetY and 0 == cropDimX and 0 == cropDimY:
return '' # user does not want to crop at all
args = ' --cropULCorner %d %d' % (offsetX, offsetY)
if cropDimX <= 0:
dimX = x - 1
else:
dimX = offsetX + cropDimX - 1
if cropDimY <= 0:
dimY = y - 1
else:
dimY = offsetY + cropDimY - 1
args += ' --cropDRCorner %d %d' % (dimX, dimY)
return args
[docs] def getUserAngle(self):
anglesDic = {0:0, 1:90, 2:180, 3:270}
return anglesDic[self.gainRot.get()]
[docs] def getUserFlip(self, imag_array):
flipDic = {0: np.asarray([[1, 0, 0], [0, 1, 0], [0, 0, 1]]),
1: np.asarray([[1, 0, 0], [0, -1, imag_array.shape[0]], [0, 0, 1]]),
2: np.asarray([[-1, 0, imag_array.shape[1]], [0, 1, 0], [0, 0, 1]])}
return flipDic[self.gainFlip.get()]
[docs] def getGPUArgs(self):
args = ' --device %(GPU)s'
if self.doLocalAlignment.get():
args += ' --processLocalShifts '
if self.skipAutotuning.get():
args += " --skipAutotuning"
args += ' --storage "%s"' % self._getExtraPath("fftBenchmark.txt")
args += ' --controlPoints %d %d %d' % (self.controlPointX, self.controlPointY, self.controlPointT)
args += ' --patchesAvg %d' % self.groupNFrames
return args
[docs] def tryProcessMovie(self, movie):
movieFolder = self._getOutputMovieFolder(movie)
x, y, n = movie.getDim()
a0, aN = self._getFrameRange(n, 'align')
s0, sN = self._getFrameRange(n, 'sum')
inputMd = os.path.join(movieFolder, 'input_movie.xmd')
writeMovieMd(movie, inputMd, a0, aN, useAlignment=False)
args = '-i "%s" ' % inputMd
args += ' -o "%s"' % self._getShiftsFile(movie)
args += ' --sampling %f' % movie.getSamplingRate()
args += ' --maxResForCorrelation %f' % self.maxResForCorrelation
args += ' --Bspline %d' % self.INTERP_MAP[self.splineOrder.get()]
if self.binFactor > 1:
args += ' --bin %f' % self.binFactor
args += self.getCropCornerArg(x, y)
args += self.getOutsideModeArg()
args += ' --frameRange %d %d' % (0, aN-a0)
args += ' --frameRangeSum %d %d' % (s0-a0, sN-a0)
args += ' --max_shift %d' % self.maxShift
if self.doSaveAveMic or self.doComputePSD:
fnAvg = self._getExtraPath(self._getOutputMicName(movie))
args += ' --oavg "%s"' % fnAvg
if self.doComputePSD:
fnInitial = os.path.join(movieFolder, "initialMic.mrc")
args += ' --oavgInitial "%s"' % fnInitial
if self.doSaveMovie:
args += ' --oaligned "%s"' % self._getExtraPath(self._getOutputMovieName(movie))
if self.inputMovies.get().getDark():
args += ' --dark "%s"' % self.inputMovies.get().getDark()
if self.inputMovies.get().getGain():
ext = pwutils.getExt(self.inputMovies.get().getFirstItem().getFileName()).lower()
if ext in ['.tif', '.tiff']:
self.flipY = True
inGainFn = self.inputMovies.get().getGain()
gainFn = xmutils.flipYImage(inGainFn, outDir = self._getExtraPath())
else:
gainFn = self.inputMovies.get().getGain()
if self.gainRot.get() != 0 or self.gainFlip.get() != 0:
gainFn = self.transformGain(gainFn)
args += ' --gain "%s"' % gainFn
if self.autoControlPoints.get():
self._setControlPoints()
if self.minLocalRes.get():
args += ' --minLocalRes %f' % self.minLocalRes
if self.useGpu.get():
args += self.getGPUArgs()
self.runJob('xmipp_cuda_movie_alignment_correlation', args, numberOfMpi=1)
else:
self.runJob('xmipp_movie_alignment_correlation', args, numberOfMpi=1)
if self.doComputePSD:
self.computePSDImages(movie, fnInitial, fnAvg)
# If the micrograph was only saved for computing the PSD
# we can remove it
pwutils.cleanPath(fnInitial)
if not self.doSaveAveMic:
pwutils.cleanPath(fnAvg)
self._saveAlignmentPlots(movie)
#--------------------------- UTILS functions ------------------------------
def _getShiftsFile(self, movie):
return self._getExtraPath(self._getMovieRoot(movie) + '_shifts.xmd')
def _stepsCheck(self):
# Input movie set can be loaded or None when checked for new inputs
# If None, we load it
self._checkNewInput()
self._checkNewOutput()
self.inputMovies.get().close()
def _setControlPoints(self):
x, y, frames = self.inputMovies.get().getDim()
Ts = self.inputMovies.get().getSamplingRate()
# one control point each 1000 A
self.controlPointX.set(max([int(x * Ts) / 1000 + 2, 3]))
self.controlPointY.set(max([int(y * Ts) / 1000 + 2, 3]))
self.controlPointT.set(max([ceil(frames/7.) + 2, 3]))
def _getMovieShifts(self, movie):
from ..convert import readShiftsMovieAlignment
""" Returns the x and y shifts for the alignment of this movie.
The shifts should refer to the original micrograph without any binning.
In case of a bining greater than 1, the shifts should be scaled.
"""
shiftsMd = md.MetaData("frameShifts@" + self._getShiftsFile(movie))
return readShiftsMovieAlignment(shiftsMd)
def _storeSummary(self, movie):
if self.doSaveAveMic and movie.hasAlignment():
s0, sN = self._getFrameRange(movie.getNumberOfFrames(), 'sum')
fstFrame, lstFrame = movie.getAlignment().getRange()
if fstFrame > s0 or lstFrame < sN:
self.summaryVar.set("Warning!!! You have selected a frame range "
"wider than the range selected to align. All "
"the frames selected without alignment "
"information, will be aligned by setting "
"alignment to 0")
def _loadMeanShifts(self, movie):
alignMd = md.MetaData("frameShifts@" + self._getShiftsFile(movie))
meanX = alignMd.getColumnValues(md.MDL_SHIFT_X)
meanY = alignMd.getColumnValues(md.MDL_SHIFT_Y)
return meanX, meanY
def _getPlotCart(self,movie):
return self._getExtraPath(self._getMovieRoot(movie)+"_plot_cart.png")
def _getPsdCorr(self,movie):
return self._getExtraPath(self._getMovieRoot(movie)+"_aligned_corrected.psd")
def _saveAlignmentPlots(self, movie):
""" Compute alignment shifts plot and save to file as a png image. """
meanX, meanY = self._loadMeanShifts(movie)
plotter = createAlignmentPlot(meanX, meanY)
plotter.savefig(self._getPlotCart(movie))
plotter.close()
def _setAlignmentInfo(self, movie, obj):
""" Set alignment info such as plot and psd filename, and
the cumulative shifts values.
Params:
movie: Pass the reference movie
obj: should pass either the created micrograph or movie
"""
obj.plotCart = Image()
obj.plotCart.setFileName(self._getPlotCart(movie))
if self.doComputePSD:
obj.psdCorr = Image()
obj.psdCorr.setFileName(self._getPsdCorr(movie))
meanX, meanY = self._loadMeanShifts(movie)
obj._xmipp_ShiftX = pwobj.CsvList()
obj._xmipp_ShiftX.set(meanX)
obj._xmipp_ShiftY = pwobj.CsvList()
obj._xmipp_ShiftY.set(meanY)
def _preprocessOutputMicrograph(self, mic, movie):
self._setAlignmentInfo(movie, mic)
def _createOutputMovie(self, movie):
alignedMovie = ProtAlignMovies._createOutputMovie(self, movie)
self._setAlignmentInfo(movie, alignedMovie)
return alignedMovie
def _validateParallelProcessing(self):
nGpus = len(self.gpuList.get().split())
nMPI = self.numberOfMpi.get()
nThreads = self.numberOfThreads.get()
errors = []
if nThreads != 1:
errors.append('Multithreading is not supported. Use a single thread or one MPI.')
validMPIs = 1 if nGpus == 1 else nGpus + 1
if self.useGpu.get() and nMPI != validMPIs:
errors.append('Invalid number of MPI. Please set it to %d, as you set %d GPU(s)' % (validMPIs, nGpus))
return errors
def _validateBinary(self):
errors = []
getXmippHome = self.getClassPackage().Plugin.getHome
cpuBinaryFn = getXmippHome('bin', 'xmipp_movie_alignment_correlation')
if self.useGpu.get() and not isXmippCudaPresent("xmipp_cuda_movie_alignment_correlation"):
errors.append('GPU version not found, make sure that Xmipp is '
'compiled with GPU\n'
'( *CUDA=True* in _scipion.conf_ + '
'_run_: $ *scipion installb xmippSrc* ).')
elif not os.path.isfile(cpuBinaryFn):
errors.append('CPU version not found for some reason, try to use the GPU version.')
return errors
def _validate(self):
if self.autoControlPoints.get():
self._setControlPoints() # make sure we work with proper values
# check execution issues
errors = ProtAlignMovies._validate(self)
errors.extend(self._validateBinary())
errors.extend(self._validateParallelProcessing())
if errors:
return errors
# check settings issues
if self.doLocalAlignment.get() and not self.useGpu.get():
errors.append("GPU is needed to do local alignment.")
if (self.controlPointX < 3):
errors.append("You have to use at least 3 control points in X dim")
if (self.controlPointY < 3):
errors.append("You have to use at least 3 control points in Y dim")
if (self.controlPointT < 3):
errors.append("You have to use at least 3 control points in T dim")
return errors
def _citations(self):
return ['strelak2020flexalign']