celutz/panorama/tasks.py

# -*- coding: utf-8 -*-
from __future__ import unicode_literals, division, print_function, absolute_import

import os
import shutil
import subprocess
import tempfile

from celery import shared_task, chain

from .gen_tiles import gen_tiles
from .hugin import parse_pto, compute_connected_components, filter_images


@shared_task
def generate_tiles(*args, **kwargs):
    return gen_tiles(*args, **kwargs)

# Pipeline of hugin-based panorama generation
# TODO: do something with the return code of the subprocess calls

@shared_task
def pto_gen(output_pto, images):
    # Projection type 1 is cylindrical
    subprocess.call(["pto_gen", "-p", "1", "-o", output_pto] + images)
    return output_pto

@shared_task
def cpfind(input_pto, output_pto):
    subprocess.call(["cpfind", "-o", output_pto, "--multirow", "--celeste", input_pto])
    return output_pto

@shared_task
def prune_images(input_pto, output_pto):
    """
    Given a pto and control points, keep the largest connected components
    of linked images.  This allows to discard images that are not
    connected to the rest of the images.

    Implementation details: the python interface to Hugin (hsi) could be
    used, but it does not export precisely the bits we need to compute the
    graph of linked images (the CPGraph class).  Thus, the simplest
    solution is to parse the .pto file by hand.
    """
    with open(input_pto) as f:
        components = compute_connected_components(parse_pto(f))
    component = max(components, key=len)
    filter_images(input_pto, component, output_pto)
    print("Keeping largest component ({} images)".format(len(component)))
    return output_pto

@shared_task
def cpclean(input_pto, output_pto):
    subprocess.call(["cpclean", "-o", output_pto, input_pto])
    return output_pto

@shared_task
def linefind(input_pto, output_pto):
    subprocess.call(["linefind", "-o", output_pto, input_pto])
    return output_pto

@shared_task
def autooptimiser(input_pto, output_pto):
    subprocess.call(["autooptimiser", "-a", "-m", "-l", "-s", "-o", output_pto, input_pto])
    return output_pto

@shared_task
def autocrop(input_pto, output_pto):
    subprocess.call(["pano_modify", "--canvas=AUTO", "--crop=AUTO", "-o", output_pto, input_pto])
    return output_pto

@shared_task
def generate_pano(input_pto, dirname, output_image):
    project_name = os.path.join(dirname, "pano")
    # hugin_executor is only available since hugin 2015.0
    # return subprocess.call(["hugin_executor", "--prefix", project_name, input_pto])
    makefile_name = os.path.join(dirname, "celery_pano.mk")
    subprocess.call(["pto2mk", "-o", makefile_name, "-p", project_name, input_pto])
    subprocess.call(["make", "-f", makefile_name])
    shutil.move(project_name + ".tif", output_image)

@shared_task
def panorama_pipeline(images, output_image):
    """Automatically assemble a panorama, using Hugin in headless mode.
    [output_image] is the name of the desired output image.  Must be a .tif
    image for now.
    """
    d = tempfile.mkdtemp(prefix='celery_hugin')
    pto = lambda filename: os.path.join(d, filename + ".pto")
    pipeline = pto_gen.s(pto("pto_gen"), images) | \
               cpfind.s(pto("cpfind")) | \
               prune_images.s(pto("prune_images")) | \
               cpclean.s(pto("cpclean")) | \
               linefind.s(pto("linefind")) | \
               autooptimiser.s(pto("autooptimiser")) | \
               autocrop.s(pto("autocrop")) | \
               generate_pano.s(d, output_image)
    return pipeline()