app.py

import os
import re
from types import SimpleNamespace
from typing import Any

import gradio as gr
import numpy as np
from detectron2 import engine

from inference import main, setup_cfg

# internal settings
NUM_PROCESSES = 1
CROP = False
SCORE_THRESHOLD = 0.8
MAX_PARTS = 5
ARGS = SimpleNamespace(
    config_file="configs/coco/instance-segmentation/swin/opd_v1_real.yaml",
    model="...",
    input_format="RGB",
    output=".output",
    cpu=True,
)


def predict(rgb_image: str, depth_image: str, intrinsics: np.ndarray, num_samples: int) -> list[Any]:
    def find_gifs(path: str) -> list[str]:
        """Scrape folders for all generated gif files."""
        for file in os.listdir(path):
            sub_path = os.path.join(path, file)
            if os.path.isdir(sub_path):
                for image_file in os.listdir(sub_path):
                    if re.match(r".*\.gif$", image_file):
                        yield os.path.join(sub_path, image_file)

    cfg = setup_cfg(ARGS)

    engine.launch(
        main,
        NUM_PROCESSES,
        args=(
            cfg,
            rgb_image,
            depth_image,
            intrinsics,
            num_samples,
            CROP,
            SCORE_THRESHOLD,
        ),
    )
    
    # process output
    # TODO: may want to select these in decreasing order of score
    pre_outputs = list(find_gifs(ARGS.output))

    outputs = []
    for idx in range(MAX_PARTS):  # hide unused components
        if idx < len(pre_outputs):
            outputs.append(gr.update(value=pre_outputs[idx], visible=True))
        else:
            outputs.append(gr.update(visible=False))
    return outputs


def variable_outputs(idx):
    idx = int(idx)


with gr.Blocks() as app:
    gr.Markdown(
        """
    # OPDMulti Demo
    Upload an image to see its range of motion.
    """
    )

    # TODO: add gr.Examples

    with gr.Row():
        rgb_image = gr.Image(
            image_mode="RGB", source="upload", type="filepath", label="RGB Image", show_label=True, interactive=True
        )
        depth_image = gr.Image(
            image_mode="L", source="upload", type="filepath", label="Depth Image", show_label=True, interactive=True
        )

    intrinsics = gr.Dataframe(
        value=[
            [
                214.85935872395834,
                0.0,
                0.0,
            ],
            [
                0.0,
                214.85935872395834,
                0.0,
            ],
            [
                125.90160319010417,
                95.13726399739583,
                1.0,
            ],
        ],
        row_count=(3, "fixed"),
        col_count=(3, "fixed"),
        datatype="number",
        type="numpy",
        label="Intrinsics matrix",
        show_label=True,
        interactive=True,
    )
    num_samples = gr.Number(
        value=10,
        label="Number of samples",
        show_label=True,
        interactive=True,
        precision=0,
        minimum=3,
        maximum=20,
    )

    submit_btn = gr.Button("Run model")

    # TODO: do we want to set a maximum limit on how many parts we render? We could also show the number of components
    # identified.
    outputs = [gr.Image(type="filepath", label=f"Part {idx + 1}", visible=False) for idx in range(MAX_PARTS)]

    # TODO: maybe need to use a queue here so we don't overload the instance
    submit_btn.click(
        fn=predict, inputs=[rgb_image, depth_image, intrinsics, num_samples], outputs=outputs, api_name="run_model"
    )

app.launch()