process_image.py

import os
import gradio as gr
from run_cmd import run_cmd
from random import randint
from PIL import Image
import tempfile

temp_path = tempfile.gettempdir()

def inference(img, size, type):
    if not img:
        raise Exception("No image!")

    _id = randint(1, 10000)
    INPUT_DIR = os.path.join(temp_path, f"input_image{str(_id)}")
    OUTPUT_DIR = os.path.join(temp_path, f"output_image{str(_id)}")
    img_in_path = os.path.join(INPUT_DIR, "input.jpg")
    img_out_path = os.path.join(OUTPUT_DIR, f"output_{size}.png")
    run_cmd(f"mkdir {INPUT_DIR}")
    run_cmd(f"mkdir {OUTPUT_DIR}")
    
    img.save(img_in_path, "PNG")

    if type == "Manga":
        run_cmd(f"python inference_manga_v2.py {img_in_path} {img_out_path}")
    else:
        run_cmd(f"python inference.py {img_in_path} {img_out_path} {type}")

    img_out = Image.open(img_out_path)

    if size == "x2":
        img_out = img_out.resize((img_out.width // 2, img_out.height // 2), resample=Image.BICUBIC)

    # Remove input and output image
    run_cmd(f"rm -rf {INPUT_DIR}")

    img_out.thumbnail((600, 600), Image.ANTIALIAS)

    return img_out, gr.File.update(value=img_out_path, visible=True)