#!/usr/bin/env python3
# coding=utf-8
"""
Calcul d'une carte de profondeur, pour l'extension « Gray Iso-Layers ».

    python depth_estimate.py image.png sortie.npy modele.onnx --model midas_small

Ce script ne tourne pas dans le Python d'Inkscape (une construction MinGW,
sans pip, ou onnxruntime ne s'installe pas) mais dans un environnement a
part, cree par `deploy.ps1 -Depth` ; l'extension le lance en sous-processus
(voir `gray_iso_layers_depth`). Dependances : onnxruntime (ou
onnxruntime-directml pour le calcul sur carte graphique), numpy, Pillow.

La sortie est un tableau numpy 2D float32 de profondeur inverse : plus la
valeur est haute, plus le point est proche.
"""

import argparse
import sys

import numpy as np

# Cote du carre presente a chaque modele.
MODELS = {"midas_small": 256, "depth_anything_v2_small": 518}

# Modeles calcules sur le processeur meme si une carte graphique est la :
# sous DirectML, MiDaS small rend une carte fausse, et son calcul est de
# toute facon plus court que l'initialisation de la carte graphique.
CPU_ONLY = {"midas_small"}

# Normalisation ImageNet, commune aux deux modeles.
_MEAN = np.array([0.485, 0.456, 0.406], dtype=np.float32)
_STD = np.array([0.229, 0.224, 0.225], dtype=np.float32)

_GPU = "DmlExecutionProvider"
_CPU = "CPUExecutionProvider"


def preprocess(picture, size):
    """Image PIL -> tenseur (1, 3, size, size) float32 attendu par le modele."""
    from PIL import Image

    rgb = picture.convert("RGB").resize((size, size), Image.BICUBIC)
    pixels = (np.asarray(rgb, dtype=np.float32) / 255.0 - _MEAN) / _STD
    return np.ascontiguousarray(pixels.transpose(2, 0, 1)[None], dtype=np.float32)


def choose_providers(available, cpu=False):
    """Fournisseurs de calcul a essayer, dans l'ordre : la carte graphique
    (DirectML) si elle est disponible et non ecartee par `cpu`, puis le
    processeur, toujours present."""
    if cpu or _GPU not in available:
        return [_CPU]
    return [_GPU, _CPU]


def _run(weights, tensor, provider):
    import onnxruntime

    options = onnxruntime.SessionOptions()
    options.log_severity_level = 3
    if provider == _GPU:
        # Exigences de DirectML.
        options.enable_mem_pattern = False
        options.execution_mode = onnxruntime.ExecutionMode.ORT_SEQUENTIAL
        # Carte designee explicitement : sans cela, l'initialisation echoue
        # sur une machine a deux cartes graphiques.
        provider = (_GPU, {"device_id": 0})
    session = onnxruntime.InferenceSession(weights, options, providers=[provider])
    name = session.get_inputs()[0].name
    # onnxruntime se replie de lui-meme sur le processeur si la carte
    # graphique ne s'initialise pas : on rend le fournisseur reellement pris.
    return session.run(None, {name: tensor})[0], session.get_providers()[0]


def estimate(picture, weights, model, cpu=False):
    """Profondeur inverse de l'image : `(tableau 2D float32, fournisseur)`.

    Le calcul se fait sur la carte graphique si le modele s'y prete (voir
    `CPU_ONLY`) ; s'il y echoue (pilote, memoire), il est refait sur le
    processeur.
    """
    import onnxruntime

    tensor = preprocess(picture, MODELS[model])
    providers = choose_providers(onnxruntime.get_available_providers(),
                                 cpu or model in CPU_ONLY)
    for provider in providers:
        try:
            depth, provider = _run(weights, tensor, provider)
            depth = np.squeeze(depth)
        except Exception:
            if provider == providers[-1]:
                raise
            continue
        if depth.ndim != 2:
            raise ValueError("sortie du modele inattendue : {}".format(depth.shape))
        return depth.astype(np.float32), provider


def main(argv=None):
    parser = argparse.ArgumentParser(description=__doc__.strip().splitlines()[0])
    parser.add_argument("image")
    parser.add_argument("output")
    parser.add_argument("weights")
    parser.add_argument("--model", choices=sorted(MODELS), default="midas_small")
    parser.add_argument("--cpu", action="store_true",
                        help="ne pas utiliser la carte graphique")
    args = parser.parse_args(argv)

    from PIL import Image

    with Image.open(args.image) as picture:
        depth, provider = estimate(picture, args.weights, args.model, args.cpu)
    with open(args.output, "wb") as handle:
        np.save(handle, depth, allow_pickle=False)
    print(provider)
    return 0


if __name__ == "__main__":
    sys.exit(main())