app.py

import gradio as gr

from chrislib.general import uninvert, view

from intrinsic.pipeline import load_models, run_pipeline

# load the intrinsic models
intrinsic_models = load_models('v2')

def generate_pipeline(models):

    def pipeline_func(image):
        return run_pipeline(models, image, device='cuda')
    
    return pipeline_func


run_pipeline = generate_pipeline(intrinsic_models)

def process_image(image):
    print(image.shape)

    # Process the image with your model pipeline
    result = run_pipeline(image)
    
    out_keys = ['hr_alb', 'dif_shd', 'pos_res']
    # Create a list to store output images
    output_images = []
    for k in out_keys:
        output_images.append(result[k])

    return output_images

interface = gr.Interface(
    fn=process_image,
    inputs=gr.inputs.Image(type="numpy"),
    outputs=gr.outputs.Image(type="numpy", label="Output Images", tool="editor"),
    live=True
)

interface.launch()