face

#!/usr/bin/env python

import argparse
import errno
import io
import os
import re
import subprocess
import sys
import json
import math
import copy

from wand.image import Image

import neural_style as ns
import face_detection as fd
import spell.client
from shutil import copy2
from git import Repo


def main():
    parser = argparse.ArgumentParser(add_help=False)
    parser.add_argument('-m', '--margin', type=int, default=0)
    parser.add_argument('-r', '--resolution', type=int, default=512)

    # args for neural style

    parser.add_argument('--style_image', type=str, default='kate.jpg')
    parser.add_argument('--main_style_weight', type=float, default=5000)
    parser.add_argument('--detail_style_weight', type=float, default=400)
    parser.add_argument('--max_iterations', type=int, default=500)
    parser.add_argument('-t', '--test_spell',
                        action='store_true', default=False)

 
    options, files = parser.parse_known_args()
    print "Face detection options ........ "
    print "Desired margin: ", options.margin
    print "Resolution: ", options.resolution
    print "Files: ", files[0]
    print "Style transfer options ........ "
    print "Style image: ", options.style_image
    print "main_style_weight: ", options.main_style_weight
    print "detail_style_weight: ", options.detail_style_weight
    print "max_iterations: ", options.max_iterations
    print "test_spell: ", options.test_spell

    if not len(files) == 1:
        print "Bailing out."
        raise ValueError("Please provide at exactly one image filename")

    filename = files[0]

    basename = os.path.splitext(os.path.basename(filename))[0]

    images_dir = os.path.join(
        os.path.dirname(os.path.realpath(__file__)),
        "image_input"
    )
    if not os.path.exists(images_dir):
        os.makedirs(images_dir)

    if options.test_spell:

        print "Prep git repo for", basename
        prep_repo(basename)

        print "Run test on spell", basename
        run_spell_test(basename)

    else:

        print "Preparing input file", filename
        prepared_filename = fd.prepare(
            filename, options.resolution, images_dir)

        print "Getting face features for ", basename
        faces = fd.get_faces(prepared_filename)

        print "Making crops for ", basename
        face_data = fd.make_crops(prepared_filename, faces, options.margin)

        print "Setting style transfer params", basename
        face_data = set_style_transfer_params(face_data, options)

        print "Set the object name", basename
        face_data["name"] = basename

        print "Write JSON for ", basename
        write_faces(face_data, prepared_filename)

        print "Prep git repo for", basename
        prep_repo(basename)

        print "Generate styled images for", basename
        runid = run_stylizer_on_spell(basename)

        print "Make image pack for Maya"
        asset_directory = collect_styled_assets(runid)
        print "Styled assets can be found in: ", asset_directory


def set_style_transfer_params(face_data, options):
    result = copy.deepcopy(face_data)

    x = result["resolution"]["x"]
    y = result["resolution"]["y"]

    result["filename_styled"] = "{}_styled{}".format( *os.path.splitext(result["filename"]))
    result["max_iterations"] = options.max_iterations
    result["style_weight"] = options.main_style_weight
    result["style_image"] = options.style_image

    for face in result["faces"]:
        face["filename_styled"] = "{}_styled{}".format(
            *os.path.splitext(face["filename"]))
        face["max_iterations"] = options.max_iterations
        face["style_weight"] = options.detail_style_weight
        face["style_image"] = options.style_image

    return result


# def do_style_transfer(images_dir, source_fn, resolution=None):
#     dest_filename = "{}_styled{}".format(*os.path.splitext(source_fn))

#     with io.open(os.path.join(images_dir, source_fn), 'rb') as image_file:
#         with Image(file=image_file) as source_image:
#             if (resolution):
#                 source_image.transform(
#                     resize="{:d}x{:d}".format(resolution["x"], resolution["y"]))
#             source_image.save(filename=os.path.join(images_dir, dest_filename))
#     return dest_filename


def write_faces(features, filename):
    json_filename = os.path.join(os.path.dirname(filename), "manifest.json")
    if features:
        with open(json_filename, 'w') as outfile:
            json.dump(features, outfile, indent=4, sort_keys=True)
    return "manifest.json"


def prep_repo(name):
    repo = Repo(".")
    repo.index.add(repo.untracked_files)
    commit = repo.git.commit('-a', '-m', "Added assets for {}".format(name))
    origin = repo.remote(name='origin')
    origin.push()
    return commit


def run_spell_test(basename):
    client = spell.client.from_environment()

    kwargs = {
        "github_url": "https://github.com/hoolymama/style_transfer.git",
        "command": "python ./test_run image_input/manifest.json",
        "machine_type": "CPU",
    }

    run = client.runs.new(**kwargs)
    print("created run: {}".format(run))
    print("run logs:")
    for line in run.logs(follow=True):
        print("\t{}".format(line))

    # print("workflow: {}".format(client.active_workflow))

    # # basic run
    # run = client.runs.new(
    #     command="echo workflow $VAR",
    #     envvars={
    #         "VAR": "SUCCESS!!!"
    #     }
    # )
    # print("created run: {}".format(run))
    # print("run logs:")
    # for line in run.logs(follow=True):
    #     print("\t{}".format(line))


def run_stylizer_on_spell(basename):
    client = spell.client.from_environment()

    kwargs = {
        "github_url": "https://github.com/hoolymama/style_transfer.git",
        "attached_resources": {"runs/4/data": "datasets"},
        "command": "python ./batch_style image_input/manifest.json",
        "machine_type": "V100",
        "framework": "tensorflow",
        "description": "style_transfer: {}".format(basename)
    }

    print "running stylizer with args: "
    print kwargs

    run = client.runs.new(**kwargs)
    print("created run: {}".format(run))
    print("run logs:")
    for line in run.logs(follow=True):
        print("\t{}".format(line))


    run.wait_status(client.runs.COMPLETE)

    runid = run.id


    out_dir = os.path.join("runs", str(runid))
    if not os.path.exists(out_dir):
        os.makedirs(out_dir)
    run.cp("image_output", out_dir)

    return runid


def collect_styled_assets(runid):
    directory = os.path.join("runs", str(runid))
    styled_directories = os.listdir(directory)

    manifest_path = "image_input/manifest.json"
    with open(manifest_path, 'r') as manifest_file:
        manifest = json.load(manifest_file)
        name = "{}_{:d}".format(manifest["name"], runid)

    asset_directory = os.path.join(directory, "{}_{:d}".format(manifest["name"], runid) )
    if not os.path.exists(asset_directory):
        os.makedirs(asset_directory)


    new_meta_filename = os.path.join(asset_directory, "meta_data.txt")
    with open(new_meta_filename, 'a+') as new_meta_file:


        for styled_name in styled_directories:
            img = os.path.join(directory, styled_name, "{}.png".format(styled_name))
            copy2(img, asset_directory)

            meta_filename = os.path.join(directory, styled_name, "meta_data.txt")
            with open(meta_filename, 'r') as meta_file:
                new_meta_file.write(meta_file.read())

    copy2(manifest_path, asset_directory)

    return asset_directory

main()