app.py

import subprocess
import os
from PIL import Image
import gradio as gr

def save_image(image, path):
    os.makedirs(os.path.dirname(path), exist_ok=True)
    image.save(path)

def load_image(path):
    return Image.open(path)

def process_image(input_image):
    input_path = "./datasets/data/test/input_image.png"
    output_dir = "./results/demo/color_pix2pix/test_latest"
    output_image_path = os.path.join(output_dir, "images", "fake_B_rgb.png")

    # Save the input image
    save_image(input_image, input_path)
    
    cmd = [
        "python", "test.py",  # Command to run the test script
        "--dataroot", "./datasets/data",  # Adjust path as needed
        "--name", "color_pix2pix",  # Model name (set according to your setup)
        "--model", "colorization",  # Model type (colorization)
        "--dataset_mode", "colorization",  # Dataset mode
        "--num_test", "1",  # Number of tests
        "--results_dir", "./results/demo",
        "--gpu_ids", "-1"  # Use CPU
    ]
    
    try:
        # Execute the command to process the image
        subprocess.run(cmd, check=True)
    except subprocess.CalledProcessError as e:
        return f"Error while running command: {e}"
    
    # Check if the output directory exists
    if not os.path.exists(output_dir):
        return f"Error: Output directory {output_dir} does not exist."

    # After processing, load the output image from the results directory
    output_files = [f for f in os.listdir(os.path.join(output_dir, "images")) if f.endswith('fake_B_rgb.png')]
    if not output_files:
        return f"Error: No output files found in {os.path.join(output_dir, 'images')}."
    
    return load_image(os.path.join(output_dir, "images", output_files[0]))

iface = gr.Interface(
    fn=process_image,
    inputs=gr.Image(type="pil"),
    outputs=gr.Image(type="pil"),
    title="Image Colorization",
    description="Upload an image to colorize it using the model."
)

if __name__ == "__main__":
    iface.launch()