Source code for gautomatch.convert

# **************************************************************************
# *
# * Authors:     J.M. De la Rosa Trevin ( [1]
# * Authors:     Grigory Sharov ( [2]
# *
# * [1] SciLifeLab, Stockholm University
# * [2] MRC Laboratory of Molecular Biology (MRC-LMB)
# *
# * 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
# * 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 ''
# *
# **************************************************************************

import os
from collections import OrderedDict

import pwem.emlib.metadata as md
from pwem.constants import NO_INDEX
from pwem.emlib.image import ImageHandler
from pwem.objects import Coordinate
from pyworkflow.object import ObjectWrap
import pyworkflow.utils as pwutils

COOR_DICT = OrderedDict([
             ("_x", md.RLN_IMAGE_COORD_X),
             ("_y", md.RLN_IMAGE_COORD_Y)

    # Additional autopicking-related metadata

[docs]def rowToObject(row, obj, attrDict, extraLabels=[]): """ This function will convert from a XmippMdRow to an EMObject. Params: row: the XmippMdRow instance (input) obj: the EMObject instance (output) attrDict: dictionary with the map between obj attributes(keys) and row MDLabels in Xmipp (values). extraLabels: a list with extra labels that could be included as properties with the label name such as: _rlnSomeThing """ obj.setEnabled(row.getValue(md.RLN_IMAGE_ENABLED, 1) > 0) for attr, label in attrDict.items(): value = row.getValue(label) if not hasattr(obj, attr): setattr(obj, attr, ObjectWrap(value)) else: getattr(obj, attr).set(value) attrLabels = attrDict.values() for label in extraLabels: if label not in attrLabels and row.hasLabel(label): labelStr = md.label2Str(label) setattr(obj, '_' + labelStr, row.getValueAsObject(label))
[docs]def rowToCoordinate(coordRow): """ Create a Coordinate from a row of a meta """ # Check that all required labels are present in the row if coordRow.containsAll(COOR_DICT): coord = Coordinate() rowToObject(coordRow, coord, COOR_DICT, extraLabels=COOR_EXTRA_LABELS) # Gautomatch starts _rlnClassNumber at 0, but relion at 1 # so let's increment its value if coord.hasAttribute('_rlnClassNumber'): coord._rlnClassNumber.increment() micName = None if coordRow.hasLabel(md.RLN_MICROGRAPH_ID): micId = int(coordRow.getValue(md.RLN_MICROGRAPH_ID)) coord.setMicId(micId) # If RLN_MICROGRAPH_NAME is not present, use the id as a name micName = micId if coordRow.hasLabel(md.RLN_MICROGRAPH_NAME): micName = coordRow.getValue(md.RLN_MICROGRAPH_NAME) coord.setMicName(micName) else: coord = None return coord
[docs]def readSetOfCoordinates(workDir, micSet, coordSet, suffix=None): """ Read from coordinates from Gautomatch .star files. For a micrograph: mic1.mrc, the expected coordinate file is: Params: workDir: where the Gautomatch output files are located. micSet: the SetOfMicrographs. coordSet: the SetOfCoordinates that will be populated. suffix: input coord file suffix """ if suffix is None: suffix = '' for mic in micSet: micBase = pwutils.removeBaseExt(mic.getFileName()) fnCoords = os.path.join(workDir, micBase + suffix) readCoordinates(mic, fnCoords, coordSet)
[docs]def readCoordinates(mic, fileName, coordsSet): if os.path.exists(fileName): for row in md.iterRows(fileName): coord = rowToCoordinate(row) coord.setX(coord.getX()) coord.setY(coord.getY()) coord.setMicrograph(mic) coordsSet.append(coord)
[docs]class CoordStarWriter: """ Helper class to write a star file containing coordinates. """ # Gautomatch cannot read default star header (with # XMIPP_STAR_1 *), # so we write directly to file HEADER = """ data_ loop_ _rlnCoordinateX #1 _rlnCoordinateY #2 _rlnAnglePsi #3 _rlnClassNumber #4 _rlnAutopickFigureOfMerit #5 """ def __init__(self, filename): self._file = open(filename, 'w') # Write header self._file.write(self.HEADER)
[docs] def writeRow(self, x, y, psi=0, classNumber=0, autopickFom=0): self._file.write("%0.6f %0.6f %0.6f %d %0.6f\n" % (x, y, psi, classNumber, autopickFom))
[docs] def close(self): self._file.close()
[docs]def writeDefectsFile(coordSet, outputFn): """ Write all coordinates in coordSet as the file as expected by Gautomatch. """ csw = CoordStarWriter(outputFn) for coord in coordSet: csw.writeRow(coord.getX(), coord.getY()) csw.close()
[docs]def writeMicCoords(mic, coordSet, outputFn): """ Write all the coordinates in coordSet as star file for micrograph mic. """ csw = CoordStarWriter(outputFn) for coord in coordSet: csw.writeRow(coord.getX(), coord.getY(), coord.getAttributeValue('_rlnAnglePsi', 0.0), coord.getAttributeValue('_rlnClassNumber', 0), coord.getAttributeValue('_rlnAutopickFigureOfMerit', 0.0)) csw.close()
[docs]def writeSetOfCoordinates(workDir, coordSet, isGlobal=False): """ Write set of coordinates from md to star file. Used only for exclusive picking. Creates .star files with bad coordinates (for each mic) and/or a single .star file with global detector defects. """ for mic in coordSet.iterMicrographs(): micBase = pwutils.removeBaseExt(mic.getFileName()) fnCoords = os.path.join(workDir, micBase + '') writeMicCoords(mic, coordSet.iterCoordinates(mic), fnCoords)
[docs]def writeSetOfCoordinatesXmipp(posDir, coordSet, ismanual=True, scale=1): """ Write a pos file on metadata format for each micrograph on the coordSet. Params: posDir: the directory where the .pos files will be written. coordSet: the SetOfCoordinates that will be read.""" boxSize = coordSet.getBoxSize() or 100 state = 'Manual' if ismanual else 'Supervised' # Create a dictionary with the pos filenames for each micrograph posDict = {} for mic in coordSet.iterMicrographs(): micIndex, micFileName = mic.getLocation() micName = os.path.basename(micFileName) if micIndex != NO_INDEX: micName = '%06d_at_%s' % (micIndex, micName) posFn = os.path.join(posDir, pwutils.replaceBaseExt(micName, "pos")) posDict[mic.getObjId()] = posFn f = None lastMicId = None c = 0 for coord in coordSet.iterItems(orderBy='_micId'): micId = coord.getMicId() if micId != lastMicId: # we need to close previous opened file if f: f.close() c = 0 f = openMd(posDict[micId], state) lastMicId = micId c += 1 if scale != 1: x = coord.getX() * scale y = coord.getY() * scale else: x = coord.getX() y = coord.getY() f.write(" %06d 1 %d %d %d %06d\n" % (coord.getObjId(), x, y, 1, micId)) if f: f.close() # Write config.xmd metadata configFn = os.path.join(posDir, 'config.xmd') writeCoordsConfig(configFn, int(boxSize), state) return posDict.values()
[docs]def writeCoordsConfig(configFn, boxSize, state): """ Write the config.xmd file needed for Xmipp picker. Params: configFn: The filename were to store the configuration. boxSize: the box size in pixels for extraction. state: picker state """ # Write config.xmd metadata print("writeCoordsConfig: state=", state) mdata = md.MetaData() # Write properties block objId = mdata.addObject() mdata.setValue(md.MDL_PICKING_PARTICLE_SIZE, int(boxSize), objId) mdata.setValue(md.MDL_PICKING_STATE, state, objId) mdata.write('properties@%s' % configFn)
[docs]def openMd(fn, state='Manual'): # We are going to write metadata directly to file to do it faster f = open(fn, 'w') ismanual = state == 'Manual' block = 'data_particles' if ismanual else 'data_particles_auto' s = """# XMIPP_STAR_1 * # data_header loop_ _pickingMicrographState %s %s loop_ _itemId _enabled _xcoor _ycoor _cost _micrographId """ % (state, block) f.write(s) return f
[docs]def writeSetOfMicrographs(micSet, filename): """ Simplified function borrowed from xmipp. """ mdata = md.MetaData() for img in micSet: objId = mdata.addObject() imgRow = md.Row() imgRow.setValue(md.MDL_ITEM_ID, int(objId)) index, fname = img.getLocation() fn = ImageHandler.locationToXmipp((index, fname)) imgRow.setValue(md.MDL_MICROGRAPH, fn) if img.isEnabled(): enabled = 1 else: enabled = -1 imgRow.setValue(md.MDL_ENABLED, enabled) imgRow.writeToMd(mdata, objId) mdata.write('Micrographs@%s' % filename)