colorization.py

import subprocess
import os
import gradio as gr
import paddlehub as hub
from pathlib import Path

# Install DeOldify model
os.system("hub install deoldify==1.0.1")

# Load the DeOldify model
model = hub.Module(name='deoldify')

# Function to colorize an image
def colorize_image(image):
    # Colorize the image using the DeOldify model
    model.predict(images=[str(image)], output_dir='./output')
    output_path = './output/' + Path(image.name).stem + ".png"
    return output_path

# Create the Gradio interface
def create_interface():
    gr.Interface(
        fn=colorize_image,
        inputs="image",
        outputs="image",
        live=True,
        title="DeOldify - Colorize Black & White Photos",
        description="Upload a black and white photo to colorize it.",
    ).launch()

if __name__ == "__main__":
    # Run the Gradio interface
    create_interface()