# **************************************************************************
# *
# * Authors: Peter Horvath [1]
# * Pablo Conesa [1]
# *
# * [1] I2PC center
# *
# * 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 3 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
import csv
import re
import emtable
import pyworkflow.utils as pwutils
import pwem.objects as emobj
from pwem.emlib.image import ImageHandler
from pwem.convert import Ccp4Header
import sphire.constants as constants
with pwutils.weakImport('tomo'):
from tomo.objects import Coordinate3D
[docs]class CoordBoxWriter:
""" Helper class to write a .BOX file for cryolo. """
def __init__(self, boxSize, yFlipHeight=None):
"""
:param boxSize: The box size of the coordinates that will be written
:param yFlipHeight: if not None, the y coordinates will be flipped
"""
self._file = None
self._boxSize = boxSize
self._halfBox = int(boxSize / 2)
self._yFlipHeight = yFlipHeight
[docs] def open(self, filename):
""" Open a new filename to write, close previous one if open. """
self.close()
self._file = open(filename, 'w')
[docs] def writeCoord(self, coord):
box = self._boxSize
half = self._halfBox
x = coord.getX() - half
if self._yFlipHeight is None:
y = coord.getY() - half
else:
y = self._yFlipHeight - coord.getY() - half
score = getattr(coord, '_cryoloScore', 0.0)
self._file.write("%s\t%s\t%s\t%s\t%s\n" % (x, y, box, box, score))
[docs] def close(self):
if self._file:
self._file.close()
[docs]class CoordBoxReader:
""" Helper class to read coordinates from .BOX files. """
def __init__(self, boxSize, yFlipHeight=None, boxSizeEstimated=False):
"""
:param boxSize: The box size of the coordinates that will be read
:param yFlipHeight: if not None, the y coordinates will be flipped
"""
self._file = None
self._boxSize = boxSize
self._halfBox = boxSize / 2.0
self._yFlipHeight = yFlipHeight
self._boxSizeEstimated = boxSizeEstimated
[docs] def open(self, filename):
""" Open a new filename to write, close previous one if open. """
self.close()
self._file = open(filename, 'r')
def _getDataLength(self):
try:
head = next(self._file) # Read the first line
cols = re.split(r'\t+', head.rstrip('\t')) # Strip it by tab
self._file.seek(0) # Return to the top of the file
return len(cols)
except StopIteration:
return None
[docs] def iterCoords(self):
reader = csv.reader(self._file, delimiter='\t')
if self._getDataLength() == 7: # crYOLO generates .cbox files with 7 columns from version > 1.5.4
for x, y, _, _, score, _, _ in reader:
# USE the imageHeight to flip or not to flip!
sciX = float(x)
sciY = float(y)
if not self._boxSizeEstimated:
sciX += self._halfBox
sciY += self._halfBox
if self._yFlipHeight is not None:
sciY = self._yFlipHeight - sciY
yield round(sciX), round(sciY), float(score)
elif self._getDataLength() == 5: # crYOLO generates .cbox files with 4 columns from version <= 1.5.4
for x, y, _, _, score in reader:
# USE the imageHeight to flip or not to flip!
sciX = round(float(x) + self._halfBox)
sciY = round(float(y) + self._halfBox)
if self._yFlipHeight is not None:
sciY = self._yFlipHeight - sciY
yield sciX, sciY, float(score)
else: # crYOLO generates .cbox files with 11 columns from version = 1.8.0
table = emtable.Table(fileName=self._file.name)
for row in table.iterRows(self._file.name, tableName='cryolo'):
x = float(row.CoordinateX)
y = float(row.CoordinateY)
score = float(row.Confidence)
if not self._boxSizeEstimated:
x += self._halfBox
y += self._halfBox
sciX = round(x)
sciY = round(y)
if self._yFlipHeight is not None:
sciY = self._yFlipHeight - sciY
yield sciX, sciY, score
[docs] def iter3DCoords(self):
table = emtable.Table(fileName=self._file.name)
for row in table.iterRows(self._file.name, tableName='cryolo'):
x = float(row.CoordinateX)
y = float(row.CoordinateY)
z = float(row.CoordinateZ)
if not self._boxSizeEstimated:
x += self._halfBox
y += self._halfBox
sciX = round(x)
sciY = round(y)
sciZ = round(z)
if self._yFlipHeight is not None:
sciY = self._yFlipHeight - sciY
yield sciX, sciY, sciZ
[docs] def close(self):
if self._file:
self._file.close()
[docs]def writeSetOfCoordinates(boxDir, coordSet, micList=None):
""" Convert a SetOfCoordinates to Cryolo box files.
Params:
boxDir: the output directory where to generate the files.
coordSet: the input SetOfCoordinates that will be converted.
micList: if not None, only coordinates from this micrographs
will be written.
"""
# Get the SOM (SetOfMicrographs)
micSet = coordSet.getMicrographs()
micIdSet = micSet.getIdSet() if micList is None else set(m.getObjId()
for m in micList)
# Get first mic from mics
mic = micSet.getFirstItem()
# Get fileName from mic
writer = CoordBoxWriter(coordSet.getBoxSize(),
getFlipYHeight(mic.getFileName()))
lastMicId = None
doWrite = True
# Loop through coordinates and generate box files
for coord in coordSet.iterItems(orderBy='_micId'):
micId = coord.getMicId()
mic = micSet[micId]
if micId != lastMicId:
doWrite = micId in micIdSet
if doWrite:
# we need to close previous opened file
writer.open(os.path.join(boxDir, getMicIdName(mic, suffix='.box')))
lastMicId = micId
if doWrite:
writer.writeCoord(coord)
writer.close()
[docs]def readSetOfCoordinates3D(tomogram, coord3DSetDict, coordsFile, boxSize,
origin=None):
reader = CoordBoxReader(boxSize)
reader.open(coordsFile)
coord3DSet = coord3DSetDict[tomogram.getObjId()]
coord3DSet.enableAppend()
coord = Coordinate3D()
for x, y, z in reader.iter3DCoords():
# Clean up objId to add as a new coordinate
coord.setObjId(None)
coord.setVolume(tomogram)
coord.setPosition(x, y, z, origin)
# Add it to the set
coord3DSet.append(coord)
reader.close()
[docs]def needToFlipOnY(filename):
""" Returns true if need to flip coordinates on Y"""
ext = pwutils.getExt(filename)
if ext in ".mrc":
header = Ccp4Header(filename, readHeader=True)
return header.getISPG() != 0 # ISPG 1, cryolo will not flip the image
return ext in constants.CRYOLO_SUPPORTED_FORMATS
[docs]def getFlipYHeight(filename):
""" Return y-Height if flipping is needed, None otherwise """
x, y, z, n = ImageHandler().getDimensions(filename)
return y if needToFlipOnY(filename) else None
[docs]def convertMicrographs(micList, micDir, prefix='mic'):
""" Convert (or simply link) input micrographs into the given directory
in a format that is compatible with crYOLO.
"""
ih = ImageHandler()
ext = pwutils.getExt(micList[0].getFileName())
def _convert(mic, newName):
ih.convert(mic, os.path.join(micDir, newName))
def _link(mic, newName):
pwutils.createAbsLink(os.path.abspath(mic.getFileName()),
os.path.join(micDir, newName))
if ext in constants.CRYOLO_SUPPORTED_FORMATS:
func = _link
else:
func = _convert
ext = '.mrc'
for mic in micList:
func(mic, getMicIdName(mic, prefix=prefix, suffix=ext))
[docs]def convertTomograms(micList, micDir):
prefix = pwutils.removeBaseExt(micList[0].getFileName())
convertMicrographs(micList, micDir, prefix=prefix)
[docs]def getMicIdName(mic, prefix='mic', suffix=''):
""" Return a name for the micrograph based on its IDs. """
return '%s%05d%s' % (prefix, mic.getObjId(), suffix)
[docs]def roundTo(number, base=1.0):
""" Returns the closest int value to number that is multiple of base"""
return int(base * round(float(number) / base))