Update app.py

app.py

@@ -2,83 +2,142 @@ 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:
-    # 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
+    img = np.transpose(img[:, :, :3], (2, 0, 1))
+    return np.expand_dims(img, axis=0).astype(np.float32)
 
 
 def post_process(img: np.array) -> np.array:
-    # 1, C, H, W -> C, H, W
     img = np.squeeze(img)
-    # C, H, W -> H, W, C
-    img = np.transpose(img, (1, 2, 0))[:, :, ::-1].astype(np.uint8)
-    return img
-
+    return np.transpose(img, (1, 2, 0))[:, :, ::-1].astype(np.uint8)
+
+# ONNX inference with session cache and GPU if available
+def get_session(model_path: str):
+    if model_path not in get_session.cache:
+        opts = onnxruntime.SessionOptions()
+        opts.intra_op_num_threads = 1
+        opts.inter_op_num_threads = 1
+        providers = ['CUDAExecutionProvider', 'CPUExecutionProvider']
+        get_session.cache[model_path] = onnxruntime.InferenceSession(model_path, opts, providers=providers)
+    return get_session.cache[model_path]
+get_session.cache = {}
 
 def inference(model_path: str, img_array: np.array) -> np.array:
-    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"
+    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"
     img = convert_pil_to_cv2(image)
     if img.ndim == 2:
         img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)
 
     if img.shape[2] == 4:
-        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()
+        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)))
+
+# === Dark Blue-Grey Theme CSS & Animations ===
+custom_css = """
+/* Dark Gradient Background */
+body .gradio-container {
+    background: linear-gradient(135deg, #0d1b2a, #1b263b, #415a77, #1b263b);
+    background-size: 400% 400%;
+    animation: bgFade 25s ease infinite;
+}
+@keyframes bgFade {
+    0% { background-position: 0% 0%; }
+    50% { background-position: 100% 100%; }
+    100% { background-position: 0% 0%; }
+}
+
+/* Title Styling */
+.fancy-title {
+    font-family: 'Poppins', sans-serif;
+    font-size: 2.8rem;
+    background: linear-gradient(90deg, #778da9, #415a77);
+    -webkit-background-clip: text;
+    -webkit-text-fill-color: transparent;
+    animation: fadeInText 2s ease-out;
+    text-align: center;
+    margin-bottom: 1rem;
+}
+@keyframes fadeInText {
+    0% { opacity: 0; transform: translateY(-10px); }
+    100% { opacity: 1; transform: translateY(0); }
+}
+
+/* Inputs & Outputs */
+.gradio-image, .gradio-gallery {
+    animation: fadeIn 1.2s ease-in;
+    border-radius: 10px;
+    box-shadow: 0 4px 12px rgba(0,0,0,0.5);
+    border: 2px solid #415a77;
+}
+@keyframes fadeIn {
+    from { opacity: 0; }
+    to { opacity: 1; }
+}
+
+/* Radio Hover */
+.gradio-radio input[type="radio"] + label:hover {
+    transform: scale(1.1);
+    color: #e0e1dd;
+    transition: transform 0.2s, color 0.2s;
+}
+
+/* Button Styling */
+.gradio-button {
+    background: linear-gradient(90deg, #1b263b, #415a77);
+    border: 1px solid #778da9;
+    border-radius: 6px;
+    color: #e0e1dd;
+    font-weight: 600;
+    padding: 10px 22px;
+    cursor: pointer;
+    box-shadow: 0 2px 6px rgba(0,0,0,0.7);
+    transition: background 0.3s, transform 0.2s;
+}
+.gradio-button:hover {
+    background: linear-gradient(90deg, #415a77, #1b263b);
+    transform: scale(1.03);
+}
+
+/* Layout tweaks */
+#upscale_btn { margin-top: 1rem; }
+.gradio-row { gap: 1rem; }
+"""
+
+# === 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")
+    btn.click(fn=upscale, inputs=[inp, model], outputs=out)
+    gr.HTML("<p style='text-align:center; color:#e0e1dd;'>Powered by ONNX Runtime & Gradio Blocks</p>")
+
+if __name__ == "__main__":
+    demo.launch()