Source code for asap.dataimport.apply_mipmaps_to_render

#!/usr/bin/env python
from functools import partial

import renderapi
from six.moves import urllib

from asap.module.render_module import StackTransitionModule
from asap.dataimport.schemas import (
    AddMipMapsToStackParameters, AddMipMapsToStackOutput)
from asap.utilities import uri_utils


if __name__ == "__main__" and __package__ is None:
    __package__ = "asap.dataimport.apply_mipmaps_to_render"

example = {
    "render": {
        "host": "em-131fs",
        "port": 8998,
        "owner": "gayathri",
        "project": "MM2",
        "client_scripts": "/allen/programs/celltypes/workgroups/em-connectomics/gayathrim/nc-em2/Janelia_Pipeline/render_20170613/render-ws-java-client/src/main/scripts"
    },
    "input_stack": "mm2_acquire_8bit",
    "output_stack": "mm2_mipmap_test",
    "mipmap_dir": "/net/aidc-isi1-prd/scratch/aibs/scratch",
    "imgformat": "tif",
    "levels": 6,
    "zstart": 1015,
    "zend": 1015
}


def addMipMapsToRender(
        render, input_stack, mipmap_prefix, imgformat, levels, z):
    resolvedtiles = render.run(
        renderapi.resolvedtiles.get_resolved_tiles_from_z,
        input_stack, z)

    for ts in resolvedtiles.tilespecs:
        mm1 = ts.ip[0]

        oldUrl = mm1.imageUrl
        filepath = urllib.parse.urlparse(oldUrl).path

        # assumes that the mipmaps are stored in the way
        #   generated by the render's mipmap client which adds the
        #   file extension in addition to the existing file extension
        #   from mipmap level 0
        if imgformat == "png":
            imgf = ".png"
        elif imgformat == "jpg":
            imgf = ".jpg"
        else:
            imgf = ".tif"

        for i in range(1, levels + 1):
            scUrl = uri_utils.uri_join(
                mipmap_prefix, str(i), "{}{}".format(
                    filepath.lstrip('/'), imgf))
            mm1 = renderapi.image_pyramid.MipMap(imageUrl=scUrl)
            ts.ip[i] = mm1
    return resolvedtiles


class AddMipMapsToStack(StackTransitionModule):
    default_schema = AddMipMapsToStackParameters
    default_output_schema = AddMipMapsToStackOutput

    def run(self):
        self.logger.debug('Applying mipmaps to stack')
        zvalues = self.get_overlapping_inputstack_zvalues()
        if len(zvalues) == 0:
            self.logger.error('No sections found for stack {}'.format(
                self.args['input_stack']))

        self.logger.debug("{}".format(zvalues))
        mypartial = partial(
            addMipMapsToRender, self.render, self.args['input_stack'],
            self.args['mipmap_prefix'], self.args['imgformat'],
            self.args['levels'])

        with renderapi.client.WithPool(self.args['pool_size']) as pool:
            allresolved = pool.map(mypartial, zvalues)

        tilespecs = [ts for ts_l in (
            resolvedtiles.tilespecs for resolvedtiles in allresolved)
                     for ts in ts_l]
        identified_tforms = list({tform.transformId: tform for tform in (
            tf for tf_l in (
                resolvedtiles.transforms for resolvedtiles in allresolved)
            for tf in tf_l)}.values())

        output_stack = (self.args['input_stack'] if
                        self.args['output_stack'] is None
                        else self.args['output_stack'])

        self.output_tilespecs_to_stack(tilespecs, output_stack,
                                       sharedTransforms=identified_tforms)

        missing_ts_zs = []
        for z in zvalues:
            job_success = self.validate_tilespecs(
                self.args['input_stack'], self.args['output_stack'], z)
            if not job_success:
                missing_ts_zs.append(z)
        self.output(
            {
                "output_stack": output_stack,
                "missing_tilespecs_zs": missing_ts_zs
            })


if __name__ == "__main__":
    mod = AddMipMapsToStack()
    mod.run()