app.py

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

def save_image(image, path):
    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.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:
        print(f"Error while running command: {e}")
        return None
    
    # Check if the output directory exists
    if not os.path.exists(output_dir):
        print(f"Error: Output directory {output_dir} does not exist.")
        return None

    # 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')]
    print(output_files)
    if not output_files:
        print(f"Error: No output files found in {os.path.join(output_dir, 'images')}.")
        return None
    
    print(output_files)

    newest_file = max(output_files, key=lambda f: os.path.getctime(os.path.join(output_dir, "images", f)))
    output_image_path = os.path.join(output_dir, "images", newest_file)
    if os.path.exists(output_image_path):
        print(f"Output image saved at {output_image_path}")
        return load_image(output_image_path)
    else:
        print(f"Error: Output image {output_image_path} not found.")
        return None

# Define the Gradio interface
iface = gr.Interface(
    fn=process_image,
    inputs=gr.Image(type="pil"),
    outputs=gr.Image(type="pil"),
    live=True,
    title="Pix2Pix Colorization",
    description="Upload an image, which will be processed using Pix2Pix model and the output will be displayed."
)


# Launch the app
iface.launch()