app.py

import gradio as gr
import numpy as np

from chrislib.general import uninvert, invert, view, view_scale

from intrinsic.pipeline import load_models, run_pipeline

intrinsic_models = load_models('v2', device='cpu')

def generate_pipeline(models):

    def pipeline_func(image, **kwargs):
        return run_pipeline(models, image, **kwargs)
    
    return pipeline_func


pipeline_func = generate_pipeline(intrinsic_models)

def process_image(image):
    print(image.shape)
    image = image.astype(np.single) / 255.

    result = pipeline_func(image, device='cpu')
    
    return [view(result['hr_alb']), 1 - invert(result['dif_shd']), view_scale(result['pos_res'])]

interface = gr.Interface(
    fn=process_image,
    inputs=gr.components.Image(type="numpy"),
    outputs=[
        gr.components.Image(type='numpy', label="Albedo"),
        gr.components.Image(type='numpy', label="Diffuse Shading"),
        gr.components.Image(type='numpy', label="Residual"),
    ],
    live=True
)

interface.launch()