Update app.py

app.py

@@ -2,138 +2,83 @@ import numpy as np
 import cv2
 import onnxruntime
 import gradio as gr
-from PIL import Image
 
-# === Upscaler Logic ===
+
 def pre_process(img: np.array) -> np.array:
-    img = np.transpose(img[:, :, :3], (2, 0, 1))
-    return np.expand_dims(img, axis=0).astype(np.float32)
+    # H, W, C -> C, H, W
+    img = np.transpose(img[:, :, 0:3], (2, 0, 1))
+    # C, H, W -> 1, C, H, W
+    img = np.expand_dims(img, axis=0).astype(np.float32)
+    return img
 
 
 def post_process(img: np.array) -> np.array:
+    # 1, C, H, W -> C, H, W
     img = np.squeeze(img)
-    return np.transpose(img, (1, 2, 0))[:, :, ::-1].astype(np.uint8)
-
-# ONNX inference with GPU if available
-def get_session(model_path: str):
-    # cache sessions
-    if model_path not in get_session.cache:
-        opts = onnxruntime.SessionOptions()
-        # multi-threading
-        opts.intra_op_num_threads = 1
-        opts.inter_op_num_threads = 1
-        # try GPU first
-        providers = ['CUDAExecutionProvider', 'CPUExecutionProvider']
-        get_session.cache[model_path] = onnxruntime.InferenceSession(
-            model_path,
-            sess_options=opts,
-            providers=providers
-        )
-    return get_session.cache[model_path]
-get_session.cache = {}
+    # C, H, W -> H, W, C
+    img = np.transpose(img, (1, 2, 0))[:, :, ::-1].astype(np.uint8)
+    return img
 
 
 def inference(model_path: str, img_array: np.array) -> np.array:
-    session = get_session(model_path)
-    inputs = {session.get_inputs()[0].name: img_array}
-    return session.run(None, inputs)[0]
-
-# PIL to BGR conversion
-def convert_pil_to_cv2(image: Image.Image) -> np.array:
-    arr = np.array(image)
-    if arr.ndim == 2:
-        return cv2.cvtColor(arr, cv2.COLOR_GRAY2BGR)
-    return arr[:, :, ::-1].copy()
-
-# Upscale handler
-def upscale(image, model_choice):
-    model_path = f"models/{model_choice}.ort"
+    options = onnxruntime.SessionOptions()
+    options.intra_op_num_threads = 1
+    options.inter_op_num_threads = 1
+    ort_session = onnxruntime.InferenceSession(model_path, options)
+    ort_inputs = {ort_session.get_inputs()[0].name: img_array}
+    ort_outs = ort_session.run(None, ort_inputs)
+
+    return ort_outs[0]
+
+
+def convert_pil_to_cv2(image):
+    # pil_image = image.convert("RGB")
+    open_cv_image = np.array(image)
+    # RGB to BGR
+    open_cv_image = open_cv_image[:, :, ::-1].copy()
+    return open_cv_image
+
+
+def upscale(image, model):
+    model_path = f"models/{model}.ort"
     img = convert_pil_to_cv2(image)
-    # handle potential alpha channel
     if img.ndim == 2:
         img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)
 
     if img.shape[2] == 4:
-        alpha = cv2.cvtColor(img[:, :, 3], cv2.COLOR_GRAY2BGR)
-        out_a = post_process(inference(model_path, pre_process(alpha)))
-        out_a = cv2.cvtColor(out_a, cv2.COLOR_BGR2GRAY)
-        rgb = img[:, :, :3]
-        out_rgb = post_process(inference(model_path, pre_process(rgb)))
-        rgba = cv2.cvtColor(out_rgb, cv2.COLOR_BGR2BGRA)
-        rgba[:, :, 3] = out_a
-        return rgba
-    return post_process(inference(model_path, pre_process(img)))
-
-# === Custom CSS for styling & animations ===
-custom_css = """
-body .gradio-container {
-    background: linear-gradient(-45deg, #ff9a9e, #fad0c4, #ffdde1);
-    background-size: 400% 400%;
-    animation: gradientBG 15s ease infinite;
-}
-@keyframes gradientBG {
-    0% { background-position: 0% 50%; }
-    50% { background-position: 100% 50%; }
-    100% { background-position: 0% 50%; }
-}
-.fancy-title {
-    font-family: 'Poppins', sans-serif;
-    font-size: 3rem;
-    background: linear-gradient(90deg, #7F7FD5, #86A8E7, #91EAE4);
-    -webkit-background-clip: text;
-    -webkit-text-fill-color: transparent;
-    animation: fadeInText 2s ease-in-out;
-    text-align: center;
-}
-@keyframes fadeInText {
-    0% { opacity: 0; transform: translateY(-20px); }
-    100% { opacity: 1; transform: translateY(0); }
-}
-.gradio-image {
-    animation: fadeIn 1s ease-in;
-    border-radius: 12px;
-    box-shadow: 0 8px 16px rgba(0,0,0,0.2);
-}
-@keyframes fadeIn {
-    from { opacity: 0; }
-    to { opacity: 1; }
-}
-.gradio-radio input[type=\"radio\"] + label:hover {
-    transform: scale(1.1);
-    transition: transform 0.2s;
-}
-.gradio-button {
-    background: linear-gradient(90deg, #FF8A00, #E52E71);
-    border: none;
-    border-radius: 8px;
-    color: white;
-    font-weight: bold;
-    padding: 12px 24px;
-    cursor: pointer;
-    transition: background 0.3s, transform 0.2s;
-}
-.gradio-button:hover {
-    background: linear-gradient(90deg, #E52E71, #FF8A00);
-    transform: scale(1.05);
-}
-#upscale_btn {
-    margin-top: 10px;
-}
-"""
-
-# === Gradio Blocks App ===
-with gr.Blocks(css=custom_css) as demo:
-    gr.HTML("<h1 class='fancy-title'>✨ Ultra AI Image Upscaler ✨</h1>")
-    with gr.Row():
-        inp = gr.Image(type="pil", label="Drop Your Image Here")
-        model = gr.Radio([
-            "modelx2", "modelx2_25JXL", "modelx4", "minecraft_modelx4"
-        ], label="Upscaler Model", value="modelx2")
-    btn = gr.Button("Upscale Image", elem_id="upscale_btn")
-    out = gr.Image(label="Upscaled Output", elem_classes="gradio-image")
-
-    btn.click(fn=upscale, inputs=[inp, model], outputs=out)
-    gr.HTML("<p style='text-align:center; color:#555;'>Powered by ONNX Runtime & Gradio Blocks</p>")
-
-if __name__ == "__main__":
-    demo.launch()
+        alpha = img[:, :, 3]  # GRAY
+        alpha = cv2.cvtColor(alpha, cv2.COLOR_GRAY2BGR)  # BGR
+        alpha_output = post_process(inference(model_path, pre_process(alpha)))  # BGR
+        alpha_output = cv2.cvtColor(alpha_output, cv2.COLOR_BGR2GRAY)  # GRAY
+
+        img = img[:, :, 0:3]  # BGR
+        image_output = post_process(inference(model_path, pre_process(img)))  # BGR
+        image_output = cv2.cvtColor(image_output, cv2.COLOR_BGR2BGRA)  # BGRA
+        image_output[:, :, 3] = alpha_output
+
+    elif img.shape[2] == 3:
+        image_output = post_process(inference(model_path, pre_process(img)))  # BGR
+
+    return image_output
+
+
+css = ".output-image, .input-image, .image-preview {height: 480px !important} "
+model_choices = ["modelx2", "modelx2 25 JXL", "modelx4", "minecraft_modelx4"]
+
+gr.Interface(
+    fn=upscale,
+    inputs=[
+        gr.Image(type="pil", label="Input Image"),
+        gr.Radio(
+            model_choices,
+            type="value",
+            value=None,
+            label="Choose Upscaler",
+        ),
+    ],
+    outputs="image",
+    title="Image Upscaler! Multiple AI",
+    description="Model: [Anchor-based Plain Net for Mobile Image Super-Resolution](https://arxiv.org/abs/2105.09750). Repository: [SR Mobile PyTorch](https://github.com/w11wo/sr_mobile_pytorch)",
+    allow_flagging="never",
+    css=css,
+).launch()