Source code for cellmaps_image_embedding.runner

#! /usr/bin/env python

import os
import sys
import time
from datetime import date
import logging
import shutil
import csv
import random
import warnings
import requests
import torch
from torch.autograd import Variable
from torch.nn import DataParallel
from torch.utils.data import DataLoader
from torch.utils.data.sampler import SequentialSampler
from tqdm import tqdm
from cellmaps_utils import constants
import cellmaps_image_embedding
from cellmaps_utils import logutils
from cellmaps_utils.provenance import ProvenanceUtil
from cellmaps_image_embedding.exceptions import CellMapsImageEmbeddingError
from cellmaps_image_embedding.dataset import *
from cellmaps_image_embedding.models import *

logger = logging.getLogger(__name__)

ABB_LABEL_INDEX = {
    "0": "Nucleoplasm",
    "1": "N. membrane",
    "2": "Nucleoli",
    "3": "N. fibrillar c.",
    "4": "N. speckles",
    "5": "N. bodies",
    "6": "ER",
    "7": "Golgi app.",
    "8": "Peroxisomes",
    "9": "Endosomes",
    "10": "Lysosomes",
    "11": "Int. fil.",
    "12": "Actin fil.",
    "13": "F. a. sites",
    "14": "Microtubules",
    "15": "M. ends",
    "16": "Cyt. bridge",
    "17": "Mitotic spindle",
    "18": "MTOC",
    "19": "Centrosome",
    "20": "Lipid droplets",
    "21": "PM",
    "22": "C. Junctions",
    "23": "Mitochondria",
    "24": "Aggresome",
    "25": "Cytosol",
    "26": "C. bodies",
    "27": "Rods & Rings"
}


[docs] class EmbeddingGenerator(object): """ Base class for implementations that generate network embeddings """ DEFAULT_FOLD = 1 DIMENSIONS = 1024 SUFFIX = '.jpg' def __init__(self, dimensions=DIMENSIONS, fold=DEFAULT_FOLD): """ Constructor """ self._dimensions = dimensions self._fold = fold self._fairscape_dataset_tuples = []
[docs] def get_dimensions(self): """ Gets number of dimensions this embedding will generate :return: number of dimensions aka vector length :rtype: int """ return self._dimensions
[docs] def get_fold(self): """ Gets fold :return: :rtype: int """ return self._fold
[docs] def get_next_embedding(self): """ Generator method for getting next embedding. Caller should implement with ``yield`` operator :raises: NotImplementedError: Subclasses should implement this :return: Embedding :rtype: list """ raise NotImplementedError('Subclasses should implement')
[docs] def get_datasets_that_need_to_be_registered(self): """ Gets any datasets that need to be registered with FAIRSCAPE :return: list of tuples in format of (dict, filepath as str) :rtype: list """ return self._fairscape_dataset_tuples
[docs] class FakeEmbeddingGenerator(EmbeddingGenerator): """ Fakes image embedding """ def __init__(self, inputdir, dimensions=EmbeddingGenerator.DIMENSIONS, fold=EmbeddingGenerator.DEFAULT_FOLD, suffix=EmbeddingGenerator.SUFFIX, img_emd_translator=None): """ Constructor :param inputdir: Directory where images reside under red, green, blue, and yellow directories :type inputdir: str :param dimensions: Desired size of output embedding :type dimensions: int :param suffix: Image suffix with starting ``.`` :type suffix: str """ super().__init__(dimensions=dimensions, fold=fold) self._inputdir = inputdir self._suffix = suffix if img_emd_translator is None: self._img_emd_translator = ImageEmbeddingFilterAndNameTranslator(image_downloaddir=inputdir, fold=fold) warnings.warn(constants.IMAGE_EMBEDDING_FILE + ' contains FAKE DATA!!!!\n' 'You have been warned\nHave a nice day\n') logger.error(constants.IMAGE_EMBEDDING_FILE + ' contains FAKE DATA!!!! ' 'You have been warned. Have a nice day') def _get_image_id_list(self): """ Looks at red directory under image directory to get a list of image ids which are the file names in that directory with last ``_`` and everything to the right of it removed from the file name :return: """ image_set = set() red_image_dir = os.path.join(self._inputdir, constants.RED) for entry in os.listdir(red_image_dir): if not entry.endswith(self._suffix): continue if not os.path.isfile(os.path.join(red_image_dir, entry)): continue # include the _ at the end cause that is also included in # image_gene_node_attributes.tsv file image_set.add(entry[: entry.rfind('_') + 1]) return list(image_set)
[docs] def get_next_embedding(self): """ Generator method for getting next embedding. Caller should implement with ``yield`` operator :raises: NotImplementedError: Subclasses should implement this :return: Embedding :rtype: list """ for image_id in self._get_image_id_list(): if image_id not in self._img_emd_translator.get_name_mapping(): continue genes = self._img_emd_translator.get_name_mapping()[image_id] for g in genes: row = [g] row.extend(np.random.normal(size=self.get_dimensions())) # sample normal distribution prob = [g] prob.extend( [random.random() for x in range(0, len(ABB_LABEL_INDEX.keys()))]) # might need to add to one yield row, prob
[docs] class DensenetEmbeddingGenerator(EmbeddingGenerator): """ Runs densenet bundled with this tool via command line to generate embedding. Why do it this way? Easier transition from the original `Densenet <https://github.com/CellProfiling/densenet>`__ code and no memory leaks """ def __init__(self, inputdir, dimensions=EmbeddingGenerator.DIMENSIONS, outdir=None, model_path=None, suffix=EmbeddingGenerator.SUFFIX, fold=EmbeddingGenerator.DEFAULT_FOLD, img_emd_translator=None): """ Constructor :param inputdir: Directory where red, blue, green, and yellow image directories reside :type inputdir: str :param dimensions: Desired size of output embedding vector :type dimensions: int :param pythonbinary: Path to python binary, if set to ``None`` the version of python that invoked this command will be used :type pythonbinary: str :param predict: Path to prediction script. Default value is the script bundled with this tool :type predict: str :param model_path: Path to model file :type model_path: str :param suffix: Image suffix with starting ``.`` :type suffix: str :param img_emd_translator: """ super().__init__(dimensions=dimensions, fold=fold) self._outdir = outdir self._inputdir = inputdir self._gpus = '' self._image_size = 1536 self._crop_size = 1024 self._device = 'cpu' self._cuda_available = False self._model_path = model_path self._suffix = suffix self._channels = 4 self._num_classes = 28 self._seeds = [0] self._augments = ['default'] self._model = None self._dataset = None self._dataloader = None if img_emd_translator is None: self._img_emd_translator = ImageEmbeddingFilterAndNameTranslator(image_downloaddir=inputdir, fold=fold) def _initialize_model(self): """ """ model = class_densenet121_large_dropout(num_classes=self._num_classes, in_channels=self._channels, pretrained=self._model_path) model = DataParallel(model) # TODO: Need to see if this is necessary # Need to properly support cpu and gpu modes model.to(self._device) # # If a node has a GPU you get an error # description of fix # https://stackoverflow.com/questions/68551032/is-there-a-way-to-use-torch-nn-dataparallel-with-cpu # and line below is the fix. model = model.module.to(self._device) model = model.eval() return model def _initialize_dataset(self): """ :return: """ dataset = ProteinDataset( self._inputdir, self._outdir, image_size=self._image_size, crop_size=self._crop_size, in_channels=self._channels, suffix=self._suffix, alt_image_ids=None) return dataset def _initialize_dataloader(self): """ :return: """ dataloader = DataLoader( self._dataset, sampler=SequentialSampler(self._dataset), batch_size=1, drop_last=False, num_workers=0, pin_memory=False, ) return dataloader def _download_model(self, model_path): """ If model_path is a URL attempt to download it to pipeline directory, otherwise return as is :param model_path: URL or file path to model file needed for image embedding :type model_path: str :return: path to model file :rtype: str """ dest_file = os.path.abspath(os.path.join(self._outdir, 'model.pth')) self._update_fairscape_dataset_tuples(dest_model=dest_file, src_url=model_path) if os.path.isfile(model_path): shutil.copy(model_path, dest_file) return dest_file with requests.get(model_path, stream=True) as r: content_size = int(r.headers.get('content-length', 0)) tqdm_bar = tqdm(desc='Downloading ' + os.path.basename(model_path), total=content_size, unit='B', unit_scale=True, unit_divisor=1024) logger.debug('Downloading ' + str(model_path) + ' of size ' + str(content_size) + 'b to ' + dest_file) try: r.raise_for_status() with open(dest_file, 'wb') as f: for chunk in r.iter_content(chunk_size=8192): f.write(chunk) tqdm_bar.update(len(chunk)) finally: tqdm_bar.close() return dest_file def _update_fairscape_dataset_tuples(self, dest_model=None, src_url=None): """ Registers model.pth file with create as a dataset """ data_dict = {'name': 'Densenet model file', 'description': 'Trained Densenet model used for classification of IF images' ' from ' + str(src_url), 'data-format': 'pth', 'keywords': ['Trained Densenet model', 'pytorch', 'classification'], 'author': cellmaps_image_embedding.__name__, 'version': cellmaps_image_embedding.__version__, 'date-published': date.today().strftime('%Y-%m-%d')} self._fairscape_dataset_tuples.append((data_dict, dest_model))
[docs] def get_next_embedding(self): """ Generator method for getting next embedding. :return: Embedding vector with 1st element :rtype: list """ self._model_path = self._download_model(self._model_path) self._model = self._initialize_model() self._dataset = self._initialize_dataset() self._dataloader = self._initialize_dataloader() for seed in self._seeds: for augment in self._augments: np.random.seed(seed) torch.manual_seed(seed) # eg. augment_default transform = eval("augment_%s" % augment) self._dataloader.dataset.set_transform(transform=transform) random_crop = (self._crop_size > 0) and (seed != 0) self._dataloader.dataset.set_random_crop(random_crop=random_crop) image_ids = np.array(self._dataloader.dataset.image_ids) for iter_index, (images, indices) in tqdm( enumerate(self._dataloader, 0), total=len(self._dataloader) ): with torch.no_grad(): if self._cuda_available: images = Variable(images.cuda()) else: images = Variable(images) logits, features = self._model(images) image_id = image_ids[iter_index] + '_' if image_id not in self._img_emd_translator.get_name_mapping(): continue genes = self._img_emd_translator.get_name_mapping()[image_id] probs = F.sigmoid(logits).cpu().data.numpy().tolist()[0] features = features.cpu().data.numpy().tolist() for g in genes: # probabilities prob_list = [g] prob_list.extend(probs) # features row = [g] row.extend(features[0]) yield row, prob_list
[docs] def get_datasets_that_need_to_be_registered(self): """ Gets model.pth dataset that needs to be registered with FAIRSCAPE. .. warning:: Must not be called before invocation of :meth:`~cellmaps_image_embedding.runner.DensenetEmbeddingGenerator.get_next_embedding` :raises CellMapsImageEmbeddingError: If this method is called before at least one invocation of :meth:`~cellmaps_image_embedding.runner.DensenetEmbeddingGenerator.get_next_embedding` :return: list of tuples in format of (dict, filepath as str) :rtype: list """ if len(self._fairscape_dataset_tuples) == 0: raise CellMapsImageEmbeddingError('get_next_embedding must be called at least ' 'once before invoking this method') return self._fairscape_dataset_tuples
[docs] class ImageEmbeddingFilterAndNameTranslator(object): """ Converts image embedding names and filters keeping only one per gene """ def __init__(self, image_downloaddir=None, fold=1): """ Constructor """ self._id_to_gene_mapping = self._gen_filtered_mapping(os.path.join(image_downloaddir, str(fold) + '_' + constants.IMAGE_GENE_NODE_ATTR_FILE)) def _gen_filtered_mapping(self, image_gene_node_attrs_file): """ Reads TSV file :param image_gene_node_attrs_file: :return: """ mapping_dict = {} with open(image_gene_node_attrs_file, 'r') as f: reader = csv.DictReader(f, delimiter='\t') for row in reader: f = row['filename'].split(',')[0] if f not in mapping_dict: mapping_dict[f] = [] mapping_dict[f].append(row['name']) if logger.isEnabledFor(logging.DEBUG): logger.debug('Mapping dict: ' + str(mapping_dict)) return mapping_dict
[docs] def get_name_mapping(self): """ Gets mapping of old name to new name :return: mapping of old name to new name :rtype: dict """ return self._id_to_gene_mapping
[docs] class CellmapsImageEmbedder(object): """ Class to run algorithm """ def __init__(self, outdir=None, inputdir=None, embedding_generator=None, skip_logging=True, name=None, organization_name=None, project_name=None, input_data_dict=None, provenance_utils=ProvenanceUtil(), provenance=None): """ Constructor :param outdir: Directory to write the results of this tool :type outdir: str :param inputdir: Directory with images that will be embedded, containing color-separated subdirectorie. Output directory from cellmaps_imagedownloader. :type inputdir: str :param embedding_generator: An object implementing the embedding generation logic, typically a subclass of :py:class:`~EmbeddingGenerator`, such as :py:class:`~DensenetEmbeddingGenerator` or :py:class:`~FakeEmbeddingGenerator`. :type embedding_generator: EmbeddingGenerator :param skip_logging: If ``True`` skip logging, if ``None`` or ``False`` do NOT skip logging :type skip_logging: bool :param name: Optional name for the dataset, used in provenance tracking. If not provided, it will be inferred from the RO-Crate in the input directory or fallback defaults. :type name: str or None :param organization_name: Optional name of the organization creating the dataset. Used for provenance metadata. :type organization_name: str or None :param project_name: Optional name of the project associated with the dataset. Used for provenance metadata. :type project_name: str or None :param input_data_dict: Dictionary of input parameters used for execution. This is captured and stored in metadata as part of reproducibility and provenance. If not provided, a default dictionary is generated. Example: .. code-block:: python {'outdir': '/output/path', 'inputdir': '/input/path'} :type input_data_dict: dict or None :param provenance_utils: Utility class instance for handling RO-Crate creation, software registration, and FAIRSCAPE dataset tracking. Default is a new :py:class:`~ProvenanceUtil` object. :type provenance_utils: ProvenanceUtil :param provenance: Optional dictionary containing explicit provenance metadata, such as dataset name, project, organization, keywords, and description. If present, it will override missing information in RO-Crate. Example: .. code-block:: python { 'name': 'Cell Image Embedding Dataset', 'organization-name': 'CM4AI', 'project-name': 'Cell Image Project', 'keywords': ['embedding', 'microscopy'], 'description': 'Embedding of IF microscopy images.' } :type provenance: dict or None """ logger.debug('In constructor') if outdir is None: raise CellMapsImageEmbeddingError('outdir is None') self._outdir = os.path.abspath(outdir) self._inputdir = os.path.abspath(inputdir) if inputdir is not None else inputdir self._start_time = int(time.time()) self._name = name self._project_name = project_name self._organization_name = organization_name self._provenance_utils = provenance_utils self._embedding_generator = embedding_generator self._softwareid = None self._input_data_dict = input_data_dict self._generated_dataset_ids = [] self._keywords = None self._description = None self._provenance = provenance self._inputdataset_ids = [] if skip_logging is None: self._skip_logging = False else: self._skip_logging = skip_logging if self._input_data_dict is None: self._input_data_dict = {'outdir': self._outdir, 'inputdir': self._inputdir, 'embedding_generator': str(self._embedding_generator), 'name': self._name, 'project_name': self._project_name, 'organization_name': self._organization_name, 'skip_logging': self._skip_logging, 'provenance': str(self._provenance) } def _update_provenance_fields(self): """ :return: """ if os.path.exists(os.path.join(self._inputdir, constants.RO_CRATE_METADATA_FILE)): prov_attrs = self._provenance_utils.get_merged_rocrate_provenance_attrs(self._inputdir, override_name=self._name, override_project_name=self._project_name, override_organization_name=self._organization_name, extra_keywords=[ 'IF Image Embedding', 'IF microscopy images', 'embedding', 'fold' + str(self._embedding_generator.get_fold())]) if self._name is None: self._name = prov_attrs.get_name() if self._organization_name is None: self._organization_name = prov_attrs.get_organization_name() if self._project_name is None: self._project_name = prov_attrs.get_project_name() self._keywords = prov_attrs.get_keywords() self._description = prov_attrs.get_description() elif self._provenance is not None: self._name = self._provenance['name'] if 'name' in self._provenance else 'Image Embedding' self._organization_name = self._provenance['organization-name'] \ if 'organization-name' in self._provenance else 'NA' self._project_name = self._provenance['project-name'] \ if 'project-name' in self._provenance else 'NA' self._keywords = self._provenance['keywords'] if 'keywords' in self._provenance else ['image'] self._description = self._provenance['description'] if 'description' in self._provenance else \ 'Embedding of Images' else: raise CellMapsImageEmbeddingError('Input directory should be an RO-Crate or provenance should be ' 'specified.') def _create_rocrate(self): """ Creates rocrate for output directory :raises CellMapsProvenanceError: If there is an error """ try: self._provenance_utils.register_rocrate(self._outdir, name=self._name, organization_name=self._organization_name, project_name=self._project_name, description=self._description, keywords=self._keywords) except TypeError as te: raise CellMapsImageEmbeddingError('Invalid provenance: ' + str(te)) except KeyError as ke: raise CellMapsImageEmbeddingError('Key missing in provenance: ' + str(ke)) def _create_output_directory(self): """ Creates output directory if it does not already exist :raises CellmapsDownloaderError: If output directory is None or if directory already exists """ if os.path.isdir(self._outdir): raise CellMapsImageEmbeddingError(self._outdir + ' already exists') os.makedirs(self._outdir, mode=0o755) for cur_color in constants.COLORS: cdir = os.path.join(self._outdir, cur_color + '_resize') if not os.path.isdir(cdir): logger.debug('Creating directory: ' + cdir) os.makedirs(cdir, mode=0o755) def _register_software(self): """ Registers this tool :raises CellMapsImageEmbeddingError: If fairscape call fails """ software_keywords = self._keywords software_keywords.extend(['tools', cellmaps_image_embedding.__name__]) software_description = self._description + ' ' + \ cellmaps_image_embedding.__description__ self._softwareid = self._provenance_utils.register_software(self._outdir, name=cellmaps_image_embedding.__name__, description=software_description, author=cellmaps_image_embedding.__author__, version=cellmaps_image_embedding.__version__, file_format='py', keywords=software_keywords, url=cellmaps_image_embedding.__repo_url__) def _register_computation(self): """ Registers computation with FAIRSCAPE """ logger.debug('Getting id of input rocrate') if os.path.exists(os.path.join(self._inputdir, constants.RO_CRATE_METADATA_FILE)): self._inputdataset_ids.append(self._provenance_utils.get_id_of_rocrate(self._inputdir)) keywords = self._keywords keywords.extend(['computation']) description = self._description + ' run of ' + cellmaps_image_embedding.__name__ self._provenance_utils.register_computation(self._outdir, name=cellmaps_image_embedding.__computation_name__, run_by=str(self._provenance_utils.get_login()), command=str(self._input_data_dict), description=description, keywords=keywords, used_software=[self._softwareid], used_dataset=self._inputdataset_ids, generated=self._generated_dataset_ids) def _register_image_embedding_file(self): """ Registers :py:const:`cellmaps_utils.constants.IMAGE_EMBEDDING_FILE` file with create as a dataset """ logger.debug('Registering embedding file with FAIRSCAPE') description = self._description description += ' file' keywords = self._keywords keywords.extend(['file']) data_dict = {'name': cellmaps_image_embedding.__name__ + ' output file', 'description': description, 'keywords': keywords, 'data-format': 'tsv', 'author': cellmaps_image_embedding.__name__, 'version': cellmaps_image_embedding.__version__, 'schema': 'https://raw.githubusercontent.com/fairscape/cm4ai-schemas/main/v0.1.0/' 'cm4ai_schema_image_embedding_emd.json', 'date-published': date.today().strftime(self._provenance_utils.get_default_date_format_str())} dset_id = self._provenance_utils.register_dataset(self._outdir, source_file=self.get_image_embedding_file(), data_dict=data_dict, skip_copy=True) self._generated_dataset_ids.append(dset_id) def _register_image_probability_file(self): """ Registers :py:const:`cellmaps_utils.constants.IMAGE_LABELS_PROBABILITY_FILE` file with create as a dataset """ logger.debug('Registering label probability file with FAIRSCAPE') description = self._description description += ' file' keywords = self._keywords keywords.extend(['file']) data_dict = {'name': cellmaps_image_embedding.__name__ + ' output file', 'description': description, 'keywords': keywords, 'data-format': 'tsv', 'author': cellmaps_image_embedding.__name__, 'version': cellmaps_image_embedding.__version__, 'schema': 'https://raw.githubusercontent.com/fairscape/cm4ai-schemas/main/v0.1.0/cm4ai_schema_image_embedding_labels_prob.json', 'date-published': date.today().strftime(self._provenance_utils.get_default_date_format_str())} dset_id = self._provenance_utils.register_dataset(self._outdir, source_file=self.get_image_probability_file(), data_dict=data_dict, skip_copy=True) self._generated_dataset_ids.append(dset_id) def _register_embedding_generator_datasets(self): """ :return: """ for dset_tuple in self._embedding_generator.get_datasets_that_need_to_be_registered(): dset_id = self._provenance_utils.register_dataset(self._outdir, source_file=dset_tuple[1], data_dict=dset_tuple[0], skip_copy=True) logger.debug('Adding embedding_generator dataset to fairscape: ' + str(dset_tuple)) self._generated_dataset_ids.append(dset_id)
[docs] def get_image_embedding_file(self): """ Gets image embedding file :return: """ return os.path.join(self._outdir, constants.IMAGE_EMBEDDING_FILE)
[docs] def get_image_probability_file(self): """ Gets image probability file :return: """ return os.path.join(self._outdir, constants.IMAGE_LABELS_PROBABILITY_FILE)
[docs] def generate_readme(self): description = getattr(cellmaps_image_embedding, '__description__', 'No description provided.') version = getattr(cellmaps_image_embedding, '__version__', '0.0.0') with open(os.path.join(os.path.dirname(__file__), 'readme_outputs.txt'), 'r') as f: readme_outputs = f.read() readme = readme_outputs.format(DESCRIPTION=description, VERSION=version) with open(os.path.join(self._outdir, 'README.txt'), 'w') as f: f.write(readme)
def _write_task_start_json(self): data = {'imagedir': self._inputdir} if self._input_data_dict is not None: data.update({'commandlineargs': self._input_data_dict}) logutils.write_task_start_json(outdir=self._outdir, start_time=self._start_time, data=data, version=cellmaps_image_embedding.__version__)
[docs] def run(self): """ Runs cellmaps_image_embedding :return: """ exitcode = 99 try: logger.debug('In run method') self._create_output_directory() if self._skip_logging is False: logutils.setup_filelogger(outdir=self._outdir, handlerprefix='cellmaps_image_embedding') self._write_task_start_json() self.generate_readme() if self._inputdir is None: raise CellMapsImageEmbeddingError('inputdir must be set') self._update_provenance_fields() self._create_rocrate() self._register_software() # generate result raw_embeddings = [] with open(self.get_image_embedding_file(), 'w', newline='') as f: with open(self.get_image_probability_file(), 'w', newline='') as pf: writer = csv.writer(f, delimiter='\t') prob_writer = csv.writer(pf, delimiter='\t') header_line = ['id'] header_line.extend([x for x in range(self._embedding_generator.get_dimensions())]) writer.writerow(header_line) header_line_prob = [''] for key in range(0, len(ABB_LABEL_INDEX.keys())): header_line_prob.append(ABB_LABEL_INDEX[str(key)]) prob_writer.writerow(header_line_prob) for row, prob_list in self._embedding_generator.get_next_embedding(): writer.writerow(row) raw_embeddings.append(row) prob_writer.writerow(prob_list) self._register_image_embedding_file() self._register_image_probability_file() self._register_embedding_generator_datasets() self._register_computation() exitcode = 0 finally: logutils.write_task_finish_json(outdir=self._outdir, start_time=self._start_time, status=exitcode) return exitcode