def correctGain(self, movieFn, outputFn, gainFn=None, darkFn=None):
"""correct a movie with both gain and dark images"""
ih = ImageHandler()
_, _, z, n = ih.getDimensions(movieFn)
numberOfFrames = max(z, n) # in case of wrong mrc stacks as volumes
def _readImgFloat(fn):
img = None
if fn:
img = ih.read(fn)
img.convert2DataType(ih.DT_FLOAT)
return img
gainImg = _readImgFloat(gainFn)
darkImg = _readImgFloat(darkFn)
img = ih.createImage()
for i in range(1, numberOfFrames + 1):
img.read((i, movieFn))
img.convert2DataType(ih.DT_FLOAT)
if darkImg:
img.inplaceSubtract(darkImg)
if gainImg:
img.inplaceMultiply(gainImg)
img.write((i, outputFn))
def _processMovie(self, movie):
allFramesSum = self._getPath('all_frames_sum.mrc')
ih = ImageHandler()
sumImg = ih.createImage()
img = ih.createImage()
n = movie.getNumberOfFrames()
fn = movie.getFileName()
sumImg.read((1, fn))
for frame in range(2, n + 1):
img.read((frame, fn))
sumImg.inplaceAdd(img)
if os.path.exists(allFramesSum):
img.read(allFramesSum)
sumImg.inplaceAdd(img)
sumImg.write(allFramesSum)
class CustomMaskDialog(ImagePreviewDialog):
def _beforePreview(self):
imgLocation = self.protocolParent.inputImage.get().getLocation()
self.dim = ImageHandler().getDimensions(imgLocation)[0]
self.lastObj = None
self.rightPreviewLabel = "Final mask"
self.message = "Generating mask..."
self.ih = ImageHandler()
self.rightImage = self.ih.createImage()
def _createPreview(self, frame):
""" Should be implemented by subclasses to
create the items preview.
"""
self._previews = []
for i, label in enumerate(MASKRESULT_LABELS):
self.previewLabel = label
previewFrame = tk.Frame(frame)
ImagePreviewDialog._createPreview(self, previewFrame)
self._previews.append(self.preview) # store all previews created
previewFrame.grid(row=i/4, column=i%4)
def _itemSelected(self, obj):
self.lastObj = obj
dialog.FlashMessage(self, self.message, func=self._computeRightPreview)
def _createVarWidgets(self, parent, varName, varLabel, row, col):
var = tk.StringVar()
self._vars[varName] = var
var.set(self.protocolParent.getAttributeValue(varName))
varLabel = tk.Label(parent, text=varLabel)
varLabel.grid(row=row, column=col*2, padx=5, pady=5)
varEntry = tk.Entry(parent, width=10, textvariable=var)
varEntry.grid(row=row, column=col*2+1, padx=5, pady=5)
def _createControls(self, frame):
self._vars = {}
inputFrame = tk.Frame(frame)
inputFrame.grid(row=0, column=0)
for i, varName in enumerate(CUSTOMMASK_VARS):
self._createVarWidgets(inputFrame, varName, CUSTOMMASK_VARS[varName], i%2, i/2)
previewBtn = HotButton(frame, text='Preview', command=self._computeRightPreview)
previewBtn.grid(row=1, column=1, padx=5, pady=5)
def getVarValue(self, varName):
return self._vars[varName].get()
def _computeRightPreview(self, e=None):
""" This function should compute the right preview
using the self.lastObj that was selected
"""
prot = self.protocolParent # short notation
tmp = prot.getProject().getTmpPath()
ext = prot.getExt()
# Convert input image to spider
imgPrefix = 'inputImage'
imgName = '%s.%s' % (imgPrefix, ext)
imgFn = os.path.join(tmp, imgName)
self.ih.convert(self.lastObj, (1, imgFn))
runCustomMaskScript(self.getVarValue('filterRadius1'),
self.getVarValue('sdFactor'),
self.getVarValue('filterRadius2'),
self.getVarValue('maskThreshold'),
workingDir=tmp, ext=ext,
inputImage=imgPrefix+'@1')
for i, preview in enumerate(self._previews):
if i == 0:
self.rightImage.read(imgFn)
else:
self.rightImage.read('%d@%s/stkmask.%s' % (i, tmp, ext))
preview.updateData(self.rightImage.getData())
class CustomMaskDialog(ImagePreviewDialog):
def _beforePreview(self):
imgLocation = self.protocolParent.inputImage.get().getLocation()
self.dim = ImageHandler().getDimensions(imgLocation)[0]
self.lastObj = None
self.rightPreviewLabel = "Final mask"
self.message = "Generating mask..."
self.ih = ImageHandler()
self.rightImage = self.ih.createImage()
def _createPreview(self, frame):
""" Should be implemented by subclasses to
create the items preview.
"""
self._previews = []
for i, label in enumerate(MASKRESULT_LABELS):
self.previewLabel = label
previewFrame = tk.Frame(frame)
ImagePreviewDialog._createPreview(self, previewFrame)
self._previews.append(self.preview) # store all previews created
previewFrame.grid(row=i / 4, column=i % 4)
def _itemSelected(self, obj):
self.lastObj = obj
dialog.FlashMessage(self, self.message, func=self._computeRightPreview)
def _createVarWidgets(self, parent, varName, varLabel, row, col):
var = tk.StringVar()
self._vars[varName] = var
var.set(self.protocolParent.getAttributeValue(varName))
varLabel = tk.Label(parent, text=varLabel)
varLabel.grid(row=row, column=col * 2, padx=5, pady=5)
varEntry = tk.Entry(parent, width=10, textvariable=var)
varEntry.grid(row=row, column=col * 2 + 1, padx=5, pady=5)
def _createControls(self, frame):
self._vars = {}
inputFrame = tk.Frame(frame)
inputFrame.grid(row=0, column=0)
for i, varName in enumerate(CUSTOMMASK_VARS):
self._createVarWidgets(inputFrame, varName,
CUSTOMMASK_VARS[varName], i % 2, i / 2)
previewBtn = HotButton(frame,
text='Preview',
command=self._computeRightPreview)
previewBtn.grid(row=1, column=1, padx=5, pady=5)
def getVarValue(self, varName):
return self._vars[varName].get()
def _computeRightPreview(self, e=None):
""" This function should compute the right preview
using the self.lastObj that was selected
"""
prot = self.protocolParent # short notation
tmp = prot.getProject().getTmpPath()
ext = prot.getExt()
# Convert input image to spider
imgPrefix = 'inputImage'
imgName = '%s.%s' % (imgPrefix, ext)
imgFn = os.path.join(tmp, imgName)
self.ih.convert(self.lastObj, (1, imgFn))
runCustomMaskScript(self.getVarValue('filterRadius1'),
self.getVarValue('sdFactor'),
self.getVarValue('filterRadius2'),
self.getVarValue('maskThreshold'),
workingDir=tmp,
ext=ext,
inputImage=imgPrefix + '@1')
for i, preview in enumerate(self._previews):
if i == 0:
self.rightImage.read(imgFn)
else:
self.rightImage.read('%d@%s/stkmask.%s' % (i, tmp, ext))
preview.updateData(self.rightImage.getData())
def createOutputStep(self):
ih = ImageHandler()
outputStack = self._getPath('particles.mrcs')
outputImg = ih.createImage()
inputParticles = self.inputParticles.get()
inputCoords = self.inputCoordinates.get()
outputSet = self._createSetOfParticles()
outputSet.copyInfo(inputParticles)
boxSize = self.boxSize.get()
b2 = int(round(boxSize / 2))
center = np.zeros((boxSize, boxSize))
ih = ImageHandler()
i = 0
outliers = 0
partIdExcluded = []
lastPartId = None
for coord in inputCoords.iterItems(
orderBy=['_subparticle._micId', '_micId', 'id']):
# The original particle id is stored in the sub-particle as micId
partId = coord._micId.get()
# Load the particle if it has changed from the last sub-particle
if partId != lastPartId:
particle = inputParticles[partId]
if particle is None:
partIdExcluded.append(partId)
self.info("WARNING: Missing particle with id %s from "
"input particles set" % partId)
else:
# Now load the particle image to extract later sub-particles
img = ih.read(particle)
x, y, _, _ = img.getDimensions()
data = img.getData()
lastPartId = partId
# If particle is not in inputParticles, subparticles will not be
# generated. Now, subtract from a subset of original particles is
# supported.
if not partId in partIdExcluded:
xpos = coord.getX()
ypos = coord.getY()
# Check that the sub-particle will not lay out of the particle
if (ypos - b2 < 0 or ypos + b2 > y or xpos - b2 < 0
or xpos + b2 > x):
outliers += 1
continue
# Crop the sub-particle data from the whole particle image
center[:, :] = data[ypos - b2:ypos + b2, xpos - b2:xpos + b2]
outputImg.setData(center)
i += 1
outputImg.write((i, outputStack))
subpart = coord._subparticle
subpart.setLocation(
(i, outputStack)) # Change path to new stack
subpart.setObjId(None) # Force to insert as a new item
outputSet.append(subpart)
if outliers:
self.info(
"WARNING: Discarded %s particles because laid out of the "
"particle (for a box size of %d" % (outliers, boxSize))
self._defineOutputs(outputParticles=outputSet)
self._defineSourceRelation(self.inputParticles, outputSet)
