texture_website.py

#
# Copyright (c) 2023-2024 Leland Brown.
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions
# are met:
#
# 1. Redistributions of source code must retain the above copyright
#    notice, this list of conditions and the following disclaimer.
#
# 2. Redistributions in binary form must reproduce the above copyright
#    notice, this list of conditions and the following disclaimer in the
#    documentation and/or other materials provided with the distribution.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
# ``AS IS'' AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
# LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
# A PARTICULAR PURPOSE ARE DISCLAIMED.  IN NO EVENT SHALL THE COPYRIGHT
# HOLDERS OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
# SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
# LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
# DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
# THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE."
#

from __future__ import annotations  # for "type | type" syntax and "tuple"

import gradio as gr
from typing import IO
import numpy as np
import pathlib
import shutil
import tempfile
from contextlib import suppress
import PIL as pillow
import PIL.Image

# import GDAL modules and set GDAL_DATA config value
from gdal_setup import gdal, osr

from terrain import TerrainData, Progress
from elevation_info import ElevationInfo
from elevation_info import TerrainFileError
from elevation_info import MissingGeoreference, MissingProjection

from constants import MIN_PREVIEW_WIDTH, MIN_PREVIEW_HEIGHT

from georeference import DistortionInfo


# default list of extensions for primary file (plus hdr.adf)
primary_extensions: list[str] = [
    '.agr', '.asc', '.bil', '.bin', '.bip', '.bsq', '.dat', '.ddf',
    '.dem', '.dt0', '.dt1', '.dt2', '.dt3', '.dt4', '.dt5', '.e00',
    '.ers', '.flt', '.grd', '.hdf', '.hgt', '.img', '.raw', '.srtm',
    '.tif', '.tiff', '.xyz',
]

# default list of file extensions always auxiliary
secondary_extensions: list[str] = [
    '.blw', '.bpw', '.bqw', '.dmw', '.hdr', '.igw', '.prj',
    '.tfw', '.tifw', '.wld',
]

# other extensions auxiliary unless no other option for primary
tertiary_extensions: list[str] = [
    '.aux', '.clr', '.dif', '.gif', '.jpeg', '.jpg', '.num', '.ovr',
    '.rrd', '.sch', '.src', '.std', '.stx', '.xml',
]

# all other extensions ending in "w" are auxiliary unless alone

# all other *.adf files are auxiliary unless no hdr.adf

# all remaining extensions may be primary

default_extensions: list[str] = (
        ['.adf'] + primary_extensions + secondary_extensions
)


def file_priority(filename: str) -> int:
    # 1 = primary file (primary_extensions and hdr.adf)
    # 2 = possibly primary (misc. files, including other *.adf)
    # 3 = possibly primary if no better option (tertiary_extensions)
    # 4 = primary only if alone (other *.*w)
    # 5 = not primary (secondary_extensions)
    path = pathlib.PurePath(filename)
    extension = path.suffix.lower()
    if extension == '.gz':
        path = pathlib.Path(path.stem)
        extension = path.suffix.lower()
    if extension in primary_extensions:
        return 1
    if extension in secondary_extensions:
        return 5
    if extension in tertiary_extensions:
        return 3
    if extension.endswith('w'):
        return 4
    if path.name.lower() == "hdr.adf":
        return 1
    return 2  # includes other *.adf files

def prioritized_files(filenames: list[str]) -> list[str]:
    filenames.sort(key=file_priority)
    return filenames

def valid_default_extension(filename: str) -> bool:
    extension = pathlib.PurePath(filename).suffix.lower()
    return extension in default_extensions

def reset():
    pass

def cleanup_temp_files(files: list[IO[bytes]]) -> None:
    for file in files:
        path = pathlib.Path(file.name)
        path.unlink(missing_ok=True)
        with suppress(OSError):  # in case dir not empty (or no longer exists)
            path.parent.rmdir()
    #temp_root = path.parent.parent
    #if temp_root.name == 'gradio':
    #    shutil.rmtree(temp_root, ignore_errors=True)

def open_terrain_file(terrain_file: list[str]) -> gdal.Dataset:
    input_files = prioritized_files(terrain_file)
    names = [pathlib.PurePath(file).name for file in input_files]

    # Cannot have two files of priority 1 - both would be primary
    if len(input_files) > 1 and file_priority(input_files[1]) == 1:
        raise TerrainFileError("More than one elevation file selected")

    first_error = None
    for index, file in enumerate(input_files):
        try:
            dataset = gdal.OpenEx(  # defaults to READONLY
                file, gdal.OF_RASTER, sibling_files=names)
            if dataset:
                # move primary file to beginning of list
                #if index > 0:
                #    input_files[1:index] = input_files[0:index - 1]
                #    input_files[0] = file
                # (then remove other files except priorities 4 and 5?
                #  and *.adf?)
                return dataset  # , input_files

            # in case of failure with no exception raised
            if not first_error:
                first_error = "Elevation data invalid or incomplete"
        except RuntimeError as e:
            if not first_error:
                # remove any reference to the directory from the error message
                #dir = str(pathlib.PurePath(file).parent) + '/'
                #first_error = "Read failed: " + str(e).replace(dir, "")

                # error messages from OpenEx are often unhelpful or misleading
                first_error = "Elevation data invalid or incomplete"

        # if file has priority 1, give up and report error
        if index == 0 and file_priority(file) == 1:
            break
    # no files were readable
    raise TerrainFileError(first_error)


class UploadStatus:
    ok: bool
    message: str

    # use keys of __annotations__ dict as slot names
    __slots__ = tuple(__annotations__)

    def __init__(self) -> None:
        self.reset()

    def reset(self) -> None:
        self.ok = False
        self.message = "No files uploaded"

    def set_ok(self) -> None:
        self.ok = True
        self.message = "SUCCESS"

    def set_error(self, err_msg: str) -> None:
        self.ok = False
        self.message = "ERROR - " + err_msg

    def status_msg(self) -> str:
        prefix = "### Status: "
        return prefix + self.message


def preview_dimensions(
        display_width: int, display_height: int,
        full_ncols: int, full_nrows: int) -> tuple[int, int, int, int]:

    width = display_width
    height = display_height

    min_width = MIN_PREVIEW_WIDTH
    min_height = MIN_PREVIEW_HEIGHT

    image_width = full_ncols
    image_height = full_nrows

    mwih = min_width * image_height
    mhiw = min_height * image_width
    if mwih >= mhiw:  # min_width is stronger constraint
        # round to nearest int; min_height never reduced here
        min_height = (2 * mwih + image_width) // (2 * image_width)
    else:  # min_height is stronger constraint
        # round to nearest int; min_width never reduced here
        min_width = (2 * mhiw + image_height) // (2 * image_height)

    wih = width * image_height
    hiw = height * image_width

    full_width = width
    full_height = height
    if wih >= hiw:  # image size limited by available height
        # round to nearest int; full_width never increased here
        full_width = (2 * hiw + image_height) // (2 * image_height)
    else:  # image size limited by available width
        # round to nearest int; full_height never increased here
        full_height = (2 * wih + image_width) // (2 * image_width)
    if full_width < min_width or full_height < min_height:
        full_width = min_width
        full_height = min_height
        if width > min_width:
            width = min_width
        elif height > min_height:
            height = min_height
    else:
        width = full_width
        height = full_height

    # scroll_vert = (full_height > height)
    # scroll_horz = (full_width > width)

    return full_width, full_height, width, height


def new_preview(
        status: UploadStatus, terrain: TerrainData,
        detail: float, contrast: float, brightness: float) \
        -> pillow.Image | None:

    # New terrain data, or detail value changed

    if not status.ok:
        return None

    # calculate fractional Laplacian
    terrain.preview.texture_shade(detail)

    # render preview image
    return update_preview(status, terrain, contrast, brightness)


#def new_preview2(
#        status: UploadStatus, terrain: TerrainData,
#        detail: float, contrast: float, brightness: float) \
#        -> pillow.Image | None:
#
#    return new_preview(status, terrain, detail, contrast, brightness)


def update_preview(
        status: UploadStatus, terrain: TerrainData,
        contrast: float, brightness: float) \
        -> pillow.Image | None:

    # Change of contrast or brightness only

    if not status.ok:
        return None

    preview_gray = terrain.preview.grayscale(contrast, brightness)

    if preview_gray is None:
        return None

    # convert to range (0.5, 255.5)
    preview_gray *= 255.0
    preview_gray += 0.5

    # set any NaN values to black
    preview_gray[np.isnan(preview_gray)] = 0.0

    # convert to integer range [0, 255]
    preview_gray = preview_gray.astype(np.uint8)

    # create Pillow Image sharing memory with preview_gray array
    preview_image = pillow.Image.fromarray(preview_gray, "L")
    #preview_image = pillow.Image.frombuffer(
    #    "L", (preview_cols, preview_rows), preview_gray, "raw", "L", 0, -1)

    width = 640
    height = 480

    full_width, full_height, width, height = preview_dimensions(
        width, height, terrain.full.ncols, terrain.full.nrows)

    # generate resized preview image
    return preview_image.resize(
        (full_width, full_height), pillow.Image.BICUBIC)


def render(
        terrain: TerrainData,
        detail: float, contrast: float, brightness: float) \
        -> tuple[str, gr.File]:

    progress: Progress = lambda n, d: None

    terrain.full.texture_shade(detail, progress)

    texture_gray = terrain.full.grayscale(
        contrast, brightness)

    # convert to range (0.5, 255.5)
    texture_gray *= 255.0
    texture_gray += 0.5

    # convert to integer range [0, 255]
    texture_gray = texture_gray.astype(np.uint8)

    driver_name = "GTiff"

    driver = gdal.GetDriverByName(driver_name)
    if not driver:
        raise RuntimeError

    has_voids = terrain.full.has_voids

    if has_voids:
        nbands = 4
        options = ["PHOTOMETRIC=RGB"]
    else:
        nbands = 1
        options = []  # default PHOTOMETRIC=MINISBLACK

    nrows, ncols = texture_gray.shape

    image_path = tempfile.NamedTemporaryFile(
        mode='wb', suffix='.tif',
        delete=False)  # delete=True, delete_on_close=False)

    try:
        output = driver.Create(
            image_path.name, ncols, nrows, nbands, gdal.GDT_Byte,
            options=options)
    except RuntimeError:
        raise

    output.SetGeoTransform(terrain.geo_transform)
    output.SetProjection(terrain.proj_wkt)

    if has_voids:
        valid = np.empty([nrows, ncols], np.uint8)
        np.isnan(terrain.full.elevation, valid)
        valid -= 1  # 0 (not NaN) -> 255 (valid), 1 (NaN) -> 0 (masked)

    if has_voids:
        for band in range(1, 4):  # RGB bands
            output.GetRasterBand(band).WriteArray(texture_gray)
        output.GetRasterBand(4).WriteArray(valid)  # alpha band
    else:
        outband = output.GetRasterBand(1)
        outband.WriteArray(texture_gray)
        # if has_voids:
        #    output.CreateMaskBand(gdal.GMF_PER_DATASET)
        #    outband.GetMaskBand().WriteArray(valid)
        del outband

    if has_voids:
        del valid

    del output  # close the output file

    #image_file = '1919-solar-eclipse.jpeg'
    return (
        image_path.name,
        gr.File(value=image_path.name, visible=True)  # False)
    )


def update_expert(expert: bool) -> gr.File:
    return gr.File(
        file_types=['file'] if expert else default_extensions)


def upload_files(
        new_files: list[IO[bytes]], expert: bool) \
        -> tuple[None, TerrainData, UploadStatus, str, list[str]]:

    valid_files: list[str] = [file.name for file in new_files]
    if not expert:
        valid_files = list(filter(valid_default_extension, valid_files))

    terrain = TerrainData()
    status = UploadStatus()

    if valid_files:
        # move file group to same temporary directory

        temp_dir = tempfile.TemporaryDirectory()
        for file in valid_files:
            # TODO: error if duplicate filenames?
            shutil.copy(file, temp_dir.name)
            # shutil.move(file, temp_dir.name)
        names = [pathlib.PurePath(file).name for file in valid_files]
        temp_dir_path = pathlib.PurePath(temp_dir.name)
        temp_files = [str(temp_dir_path.joinpath(name)) for name in names]

        # load terrain data

        dataset: gdal.Dataset | None
        try:
            dataset = open_terrain_file(temp_files)
            elev = ElevationInfo(dataset)
            dataset = None  # close input file
            distortion: DistortionInfo = elev.distortion()
            terrain = TerrainData(elev, distortion)
            del distortion
            del elev
            if terrain.full.has_voids:
                raise RuntimeError(
                    "File contains void pixels (NODATA)\n"
                    "This version of Texture Shader cannot process terrain "
                    "regions containing void pixels.\n"
                    "Choose a file with a rectangular area "
                    "of complete data."
                )
            status.set_ok()
        except RuntimeError as e:
            status.set_error(str(e))
        finally:
            dataset = None  # close input file

    return None, terrain, status, status.status_msg(), valid_files


title_text: str = (
"""
# <p style="text-align: center;">Online Texture Shader</p>

This page converts a terrain elevation data file to a "texture shaded" map image.
See [TextureShading.com](https://www.textureshading.com/) for more information.
"""
)

step1_text: str = (
"""
## Step 1: Upload Terrain Data
- Select multiple files when appropriate - include any related auxiliary files
with extensions .hdr, .prj, .tfw, etc.
"""
)

step2_text: str = (
"""
## Step 2: Adjust Preview Image to Your Liking
- Use the sliders below the image to make changes
"""
)

step3_text: str = (
"""
## Step 3: When Ready, Generate Full Image for Download
"""
)

privacy_text: str = (
"""
<br>
Privacy Policy: No cookies are used, and no user data is retained.
All uploaded files are used only to generate your output and then are deleted from the server.
"""
)


if __name__ == '__main__':
    gdal.UseExceptions()  # also suppresses error messages to stdout
    gdal.SetConfigOption("GDAL_TIFF_INTERNAL_MASK", "YES")  # for output files

    #demo = gr.Interface(fn=greet, inputs="text", outputs="text")
    #demo = gr.Interface(sepia, gr.Image(shape=(200, 200)), "image")

    with gr.Blocks() as demo:
        status = gr.State(value=UploadStatus())
        terrain = gr.State(value=TerrainData())

        gr.Markdown(title_text)
        gr.Markdown(step1_text)
        with gr.Group():
            with gr.Row():
                with gr.Column():
                    terrain_file = gr.File(
                        label="Terrain Elevation File(s)",
                        file_count='multiple',
                        file_types=default_extensions,
                        height=240)
                    #terrain_file = gr.UploadButton(
                    #    label="Click to Upload Terrain Elevation File(s)",
                    #    file_count='multiple',
                    #    file_types=default_extensions)
                    expert = gr.Checkbox(
                        value=False,
                        label="Accept uncommon file types (\"Trust me, I know what I'm doing...\")",
                        info=None)
                with gr.Column():
                    status_msg = gr.Markdown(UploadStatus().status_msg())
                    file_list = gr.File(
                        label="Files Uploaded:",
                        file_count='multiple'
                    )

        gr.Markdown(step2_text)
        with gr.Group():
            preview_img = gr.Image(
                label="Preview Image", image_mode='L',  # height=480,
                show_download_button=False
            )
            detail = gr.Slider(
                minimum=0, maximum=200, value=50, step=5, label="Texture Detail")
            contrast = gr.Slider(
                minimum=-30, maximum=30, value=0, step=1,
                label="Midtone Contrast * (usually needs adjustment)")
            brightness = gr.Slider(
                minimum=-30, maximum=30, value=0, step=1, label="Midtone Brightness")

        gr.Markdown(step3_text)
        with gr.Group():
            with gr.Row():
                gr.Markdown("")
                go = gr.Button(value="Render Full Image")
                gr.Markdown("")
            output_img = gr.Image(
                label="Output Image", image_mode='L',
                visible=False
            )
            with gr.Row():
                gr.Markdown("")
                download = gr.File(label="Download Here:", visible=False)
                gr.Markdown("")


        gr.Markdown(privacy_text)

        expert.change(update_expert, inputs=expert, outputs=terrain_file)

        #terrain_file.clear(reset)  # only with File, not UploadButton
        terrain_file.upload(
            upload_files,
            inputs=[terrain_file, expert],
            outputs=[terrain_file, terrain, status, status_msg, file_list]
        ).then(
            fn=new_preview,
            inputs=[status, terrain, detail, contrast, brightness],
            outputs=preview_img
        )

        detail.change(
            fn=new_preview,
            inputs=[status, terrain, detail, contrast, brightness],
            outputs=preview_img
        )

        gr.on(
            triggers=[contrast.change, brightness.change],
            fn=update_preview,
            inputs=[status, terrain, contrast, brightness],
            outputs=preview_img
        )

        go.click(
            render, inputs=[terrain, detail, contrast, brightness],
            outputs=[output_img, download])

    demo.launch(share=True)