OpenDroneMap-WebODM/coreplugins/cesiumion/api_views.py

import sys
import time
import logging
import requests
from os import path
from enum import Enum
from itertools import chain as iter_chain

from app.plugins.views import TaskView
from app.plugins.worker import run_function_async
from app.plugins.data_store import GlobalDataStore
from app.plugins import signals as plugin_signals

from worker.celery import app
from django.dispatch import receiver
from django.utils.translation import ugettext_lazy as _
from rest_framework.fields import ChoiceField, CharField, JSONField
from rest_framework.views import APIView
from rest_framework.response import Response
from rest_framework import status
from rest_framework import serializers

from .globals import PROJECT_NAME, ION_API_URL
from .uploader import upload_to_ion


pluck = lambda dic, *keys: [dic[k] if k in dic else None for k in keys]


###                        ###
#         API UTILS          #
###                        ###
def get_key_for(task_id, key):
    return "task_{}_{}".format(str(task_id), key)


def del_asset_info(task_id, asset_type, ds=None):
    if ds is None:
        ds = GlobalDataStore(PROJECT_NAME)
    ds.del_key(get_key_for(task_id, asset_type.value))


def set_asset_info(task_id, asset_type, json, ds=None):
    if ds is None:
        ds = GlobalDataStore(PROJECT_NAME)
    return ds.set_json(get_key_for(task_id, asset_type.value), json)


def get_asset_info(task_id, asset_type, default=None, ds=None):
    if default is None:
        default = {
            "id": None,
            "upload": {"progress": 0, "active": False},
            "process": {"progress": 0, "active": False},
            "error": "",
        }
    if ds is None:
        ds = GlobalDataStore(PROJECT_NAME)
    return ds.get_json(get_key_for(task_id, asset_type.value), default)


def is_asset_task(asset_meta):
    is_error = len(asset_meta["error"]) > 0
    return asset_meta["upload"]["active"] or asset_meta["process"]["active"] or is_error


def get_processing_assets(task_id):
    ispc = app.control.inspect()
    ion_tasks = set()
    active = set()
    from uuid import UUID

    for wtask in iter_chain(*ispc.active().values(), *ispc.reserved().values()):
        args = eval(wtask["args"])
        if len(args) < 2:
            continue
        ion_tasks.add((str(args[0]), AssetType[args[1]]))

    for asset_type in AssetType:
        asset_info = get_asset_info(task_id, asset_type)
        ion_task_id = (task_id, asset_type)
        if not is_asset_task(asset_info) or ion_task_id in ion_tasks:
            continue
        active.add(asset_type)

    return active


###                        ###
#      MODEL CONFIG          #
###                        ###
class AssetType(str, Enum):
    ORTHOPHOTO = "ORTHOPHOTO"
    TERRAIN_MODEL = "TERRAIN_MODEL"
    SURFACE_MODEL = "SURFACE_MODEL"
    POINTCLOUD = "POINTCLOUD"
    TEXTURED_MODEL = "TEXTURED_MODEL"


class SourceType(str, Enum):
    RASTER_IMAGERY = "RASTER_IMAGERY"
    RASTER_TERRAIN = "RASTER_TERRAIN"
    TERRAIN_DATABASE = "TERRAIN_DATABASE"
    CITYGML = "CITYGML"
    KML = "KML"
    CAPTURE = "3D_CAPTURE"
    MODEL = "3D_MODEL"
    POINTCLOUD = "POINT_CLOUD"


class OutputType(str, Enum):
    IMAGERY = "IMAGERY"
    TILES = "3DTILES"
    TERRAIN = "TERRAIN"


ASSET_TO_FILE = {
    AssetType.ORTHOPHOTO: "orthophoto.tif",
    AssetType.TERRAIN_MODEL: "dtm.tif",
    AssetType.SURFACE_MODEL: "dsm.tif",
    AssetType.POINTCLOUD: "georeferenced_model.laz",
    AssetType.TEXTURED_MODEL: "textured_model.zip",
}

FILE_TO_ASSET = dict([reversed(i) for i in ASSET_TO_FILE.items()])

ASSET_TO_OUTPUT = {
    AssetType.ORTHOPHOTO: OutputType.IMAGERY,
    AssetType.TERRAIN_MODEL: OutputType.TERRAIN,
    AssetType.SURFACE_MODEL: OutputType.TERRAIN,
    AssetType.POINTCLOUD: OutputType.TILES,
    AssetType.TEXTURED_MODEL: OutputType.TILES,
}

ASSET_TO_SOURCE = {
    AssetType.ORTHOPHOTO: SourceType.RASTER_IMAGERY,
    AssetType.TERRAIN_MODEL: SourceType.RASTER_TERRAIN,
    AssetType.SURFACE_MODEL: SourceType.RASTER_TERRAIN,
    AssetType.POINTCLOUD: SourceType.POINTCLOUD,
    AssetType.TEXTURED_MODEL: SourceType.CAPTURE,
}

###                        ###
#         RECIEVERS          #
###                        ###
@receiver(plugin_signals.task_removed, dispatch_uid="oam_on_task_removed")
@receiver(plugin_signals.task_completed, dispatch_uid="oam_on_task_completed")
def oam_cleanup(sender, task_id, **kwargs):
    # When a task is removed, simply remove clutter
    # When a task is re-processed, make sure we can re-share it if we shared a task previously
    for asset_type in AssetType:
        del_asset_info(task_id, asset_type)


###                        ###
#         API VIEWS          #
###                        ###
class EnumField(ChoiceField):
    default_error_messages = {"invalid": _("No matching enum type.")}

    def __init__(self, **kwargs):
        self.enum_type = kwargs.pop("enum_type")
        choices = [enum_item.value for enum_item in self.enum_type]
        self.choice_set = set(choices)
        super().__init__(choices, **kwargs)

    def to_internal_value(self, data):
        if data in self.choice_set:
            return self.enum_type[data]
        self.fail("invalid")

    def to_representation(self, value):
        if not value:
            return None
        return value.value


class UploadSerializer(serializers.Serializer):
    token = CharField()
    name = CharField()
    asset_type = EnumField(enum_type=AssetType)
    description = CharField(default="", required=False, allow_blank=True)
    attribution = CharField(default="", required=False, allow_blank=True)
    options = JSONField(default={}, required=False)


class UpdateIonAssets(serializers.Serializer):
    token = CharField()


class ShareTaskView(TaskView):
    def get(self, request, pk=None):
        task = self.get_and_check_task(request, pk)

        assets = []
        for file_name in task.available_assets:
            if file_name not in FILE_TO_ASSET:
                continue
            asset_type = FILE_TO_ASSET[file_name]

            asset_info = get_asset_info(task.id, asset_type)
            ion_id = asset_info["id"]
            is_error = len(asset_info["error"]) > 0
            is_task = is_asset_task(asset_info)
            is_exported = asset_info["id"] is not None and not is_task

            assets.append(
                {
                    "type": asset_type,
                    "isError": is_error,
                    "isTask": is_task,
                    "isExported": is_exported,
                    **asset_info,
                }
            )

        return Response({"items": assets}, status=status.HTTP_200_OK)

    def post(self, request, pk=None):
        task = self.get_and_check_task(request, pk)
        serializer = UploadSerializer(data=request.data)
        serializer.is_valid(raise_exception=True)

        token, asset_type, name, description, attribution, options = pluck(
            serializer.validated_data,
            "token",
            "asset_type",
            "name",
            "description",
            "attribution",
            "options",
        )
        asset_path = task.get_asset_download_path(ASSET_TO_FILE[asset_type])

        # Skip already processing tasks
        if asset_type not in get_processing_assets(task.id):
            if asset_type == AssetType.TEXTURED_MODEL and "position" not in options:
                extent = None
                if task.dsm_extent is not None:
                    extent = task.dsm_extent.extent
                if task.dtm_extent is not None:
                    extent = task.dtm_extent.extent
                if extent is None:
                    print(f"Unable to find task boundary: {task}")
                else:
                    lng, lat = extent[0], extent[1]
                    # height is set to zero as model is already offset
                    options["position"] = [lng, lat, 0]

            del_asset_info(task.id, asset_type)
            asset_info = get_asset_info(task.id, asset_type)
            asset_info["upload"]["active"] = True
            set_asset_info(task.id, asset_type, asset_info)

            run_function_async(upload_to_ion,
                task.id,
                asset_type,
                token,
                asset_path,
                name,
                description,
                attribution,
                options,
            )
        else:
            print(f"Ignore running ion task {task.id} {str(asset_type)}")

        return Response(status=status.HTTP_200_OK)


class RefreshIonTaskView(TaskView):
    def post(self, request, pk=None):
        serializer = UpdateIonAssets(data=request.data)
        serializer.is_valid(raise_exception=True)

        task = self.get_and_check_task(request, pk)
        token = serializer.validated_data["token"]
        headers = {"Authorization": f"Bearer {token}"}

        is_updated = False
        # ion cleanup check
        for asset_type in AssetType:
            asset_info = get_asset_info(task.id, asset_type)
            ion_id = asset_info["id"]
            if ion_id is None:
                continue
            res = requests.get(f"{ION_API_URL}/assets/{ion_id}", headers=headers)
            if res.status_code != 200:
                del_asset_info(task.id, asset_type)
                is_updated = True

        # dead task cleanup
        for asset_type in get_processing_assets(task.id):
            del_asset_info(task.id, asset_type)
            is_updated = True

        return Response({"updated": is_updated}, status=status.HTTP_200_OK)


class ClearErrorsTaskView(TaskView):
    def post(self, request, pk=None):
        task = self.get_and_check_task(request, pk)
        for asset_type in AssetType:
            asset_info = get_asset_info(task.id, asset_type)
            if len(asset_info["error"]) <= 0:
                continue
            del_asset_info(task.id, asset_type)

        return Response({"complete": True}, status=status.HTTP_200_OK)