Update app.py

app.py

@@ -8,33 +8,45 @@ import cv2
 import gradio as gr
 import numpy as np
 import torch
+
 from basicsr.archs.rrdbnet_arch import RRDBNet
 from realesrgan import RealESRGANer
 
+
 # =========================
 # CONFIG
 # =========================
 OUTSCALE = 2
 
-# Peso oficial do anime model mostrado no README do Real-ESRGAN
-MODEL_URL = "https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.2.4/RealESRGAN_x4plus_anime_6B.pth"
+MODEL_URL = (
+    "https://github.com/xinntao/Real-ESRGAN/releases/download/"
+    "v0.2.2.4/RealESRGAN_x4plus_anime_6B.pth"
+)
+
 MODEL_DIR = Path("weights")
 MODEL_PATH = MODEL_DIR / "RealESRGAN_x4plus_anime_6B.pth"
 
 
-def ensure_model() -> str:
+# =========================
+# DOWNLOAD MODEL
+# =========================
+def ensure_model():
     MODEL_DIR.mkdir(parents=True, exist_ok=True)
 
-    if MODEL_PATH.exists() and MODEL_PATH.stat().st_size > 0:
+    if MODEL_PATH.exists():
         return str(MODEL_PATH)
 
+    print("Downloading model...")
     urlretrieve(MODEL_URL, MODEL_PATH)
+
     return str(MODEL_PATH)
 
 
+# =========================
+# BUILD UPSAMPLER
+# =========================
 def build_upsampler():
     device = "cuda" if torch.cuda.is_available() else "cpu"
-    use_half = device == "cuda"
 
     model_path = ensure_model()
 
@@ -47,117 +59,241 @@ def build_upsampler():
         scale=4,
     )
 
-    return RealESRGANer(
+    upsampler = RealESRGANer(
         scale=4,
         model_path=model_path,
         model=model,
         tile=256,
         tile_pad=10,
         pre_pad=0,
-        half=use_half,
+        half=torch.cuda.is_available(),
         device=device,
     )
 
+    return upsampler
+
 
 UPSAMPLER = build_upsampler()
 
 
-def upscale_one_image(image: np.ndarray) -> np.ndarray:
-    if image is None:
+# =========================
+# UPSCALE SINGLE IMAGE
+# =========================
+def upscale_image(img):
+
+    if img is None:
         return None
 
-    if image.dtype != np.uint8:
-        image = np.clip(image, 0, 1)
-        image = (image * 255).astype(np.uint8)
+    if img.dtype != np.uint8:
+        img = (np.clip(img, 0, 1) * 255).astype(np.uint8)
 
-    if image.ndim == 3 and image.shape[2] == 4:
-        rgb = image[:, :, :3]
-        alpha = image[:, :, 3]
+    has_alpha = (
+        len(img.shape) == 3
+        and img.shape[2] == 4
+    )
+
+    if has_alpha:
+        rgb = img[:, :, :3]
+        alpha = img[:, :, 3]
 
         bgr = cv2.cvtColor(rgb, cv2.COLOR_RGB2BGR)
-        out_bgr, _ = UPSAMPLER.enhance(bgr, outscale=OUTSCALE)
-        out_rgb = cv2.cvtColor(out_bgr, cv2.COLOR_BGR2RGB)
 
-        alpha_up = cv2.resize(
+        output, _ = UPSAMPLER.enhance(
+            bgr,
+            outscale=OUTSCALE
+        )
+
+        output = cv2.cvtColor(
+            output,
+            cv2.COLOR_BGR2RGB
+        )
+
+        alpha = cv2.resize(
             alpha,
-            (alpha.shape[1] * OUTSCALE, alpha.shape[0] * OUTSCALE),
-            interpolation=cv2.INTER_CUBIC,
+            (
+                alpha.shape[1] * OUTSCALE,
+                alpha.shape[0] * OUTSCALE
+            ),
+            interpolation=cv2.INTER_CUBIC
         )
-        return np.dstack([out_rgb, alpha_up])
 
-    bgr = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
-    out_bgr, _ = UPSAMPLER.enhance(bgr, outscale=OUTSCALE)
-    out_rgb = cv2.cvtColor(out_bgr, cv2.COLOR_BGR2RGB)
-    return out_rgb
+        output = np.dstack([output, alpha])
+
+        return output
 
+    else:
+        bgr = cv2.cvtColor(
+            img,
+            cv2.COLOR_RGB2BGR
+        )
+
+        output, _ = UPSAMPLER.enhance(
+            bgr,
+            outscale=OUTSCALE
+        )
 
+        output = cv2.cvtColor(
+            output,
+            cv2.COLOR_BGR2RGB
+        )
+
+        return output
+
+
+# =========================
+# BATCH PROCESS
+# =========================
 def process_batch(files):
+
     if not files:
         return [], None
 
     previews = []
 
     with tempfile.TemporaryDirectory() as tmpdir:
+
         tmpdir = Path(tmpdir)
-        out_dir = tmpdir / "upscaled"
-        out_dir.mkdir(parents=True, exist_ok=True)
 
-        for file_obj in files:
-            in_path = Path(file_obj.name)
-            image = cv2.imread(str(in_path), cv2.IMREAD_UNCHANGED)
+        output_dir = tmpdir / "outputs"
+        output_dir.mkdir(
+            parents=True,
+            exist_ok=True
+        )
+
+        for file in files:
+
+            input_path = Path(file.name)
+
+            img = cv2.imread(
+                str(input_path),
+                cv2.IMREAD_UNCHANGED
+            )
 
-            if image is None:
+            if img is None:
                 continue
 
-            if image.ndim == 3 and image.shape[2] == 4:
-                image = cv2.cvtColor(image, cv2.COLOR_BGRA2RGBA)
-            elif image.ndim == 3:
-                image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+            # BGR -> RGB
+            if len(img.shape) == 3:
+
+                if img.shape[2] == 4:
+                    img = cv2.cvtColor(
+                        img,
+                        cv2.COLOR_BGRA2RGBA
+                    )
+                else:
+                    img = cv2.cvtColor(
+                        img,
+                        cv2.COLOR_BGR2RGB
+                    )
+
+            result = upscale_image(img)
 
-            result = upscale_one_image(image)
             if result is None:
                 continue
 
-            out_name = f"{in_path.stem}_2x.png"
-            out_path = out_dir / out_name
+            out_name = f"{input_path.stem}_2x.png"
+
+            out_path = output_dir / out_name
+
+            # RGB -> BGR
+            if len(result.shape) == 3:
+
+                if result.shape[2] == 4:
+                    save_img = cv2.cvtColor(
+                        result,
+                        cv2.COLOR_RGBA2BGRA
+                    )
+                else:
+                    save_img = cv2.cvtColor(
+                        result,
+                        cv2.COLOR_RGB2BGR
+                    )
 
-            if result.ndim == 3 and result.shape[2] == 4:
-                save_img = cv2.cvtColor(result, cv2.COLOR_RGBA2BGRA)
             else:
-                save_img = cv2.cvtColor(result, cv2.COLOR_RGB2BGR)
+                save_img = result
+
+            cv2.imwrite(
+                str(out_path),
+                save_img
+            )
 
-            cv2.imwrite(str(out_path), save_img)
-            previews.append((result, out_name))
+            previews.append(
+                (result, out_name)
+            )
 
         zip_path = tmpdir / "upscaled_images.zip"
-        with zipfile.ZipFile(zip_path, "w", compression=zipfile.ZIP_DEFLATED) as zf:
-            for img_file in out_dir.iterdir():
-                zf.write(img_file, arcname=img_file.name)
 
-        final_zip = Path(tempfile.gettempdir()) / "upscaled_images.zip"
-        final_zip.write_bytes(zip_path.read_bytes())
+        with zipfile.ZipFile(
+            zip_path,
+            "w",
+            compression=zipfile.ZIP_DEFLATED
+        ) as zipf:
+
+            for img_file in output_dir.iterdir():
+
+                zipf.write(
+                    img_file,
+                    arcname=img_file.name
+                )
+
+        final_zip = (
+            Path(tempfile.gettempdir())
+            / "upscaled_images.zip"
+        )
+
+        final_zip.write_bytes(
+            zip_path.read_bytes()
+        )
 
         return previews, str(final_zip)
 
 
+# =========================
+# UI
+# =========================
 with gr.Blocks() as demo:
-    gr.Markdown("# Anime Upscaler 2x\nUpload em lote e baixe um ZIP.")
 
-    files_in = gr.Files(
-        label="Envie várias imagens",
+    gr.Markdown(
+        "# Anime Upscaler 2x\n"
+        "Batch upscale com ZIP."
+    )
+
+    files_in = gr.File(
+        label="Envie imagens",
         file_types=["image"],
         file_count="multiple",
     )
 
-    run_btn = gr.Button("Processar")
-    gallery_out = gr.Gallery(label="Prévia", columns=2, height=420)
-    zip_out = gr.File(label="Baixar ZIP")
+    run_btn = gr.Button(
+        "Processar"
+    )
+
+    gallery_out = gr.Gallery(
+        label="Preview",
+        columns=2,
+        height=400
+    )
+
+    zip_out = gr.File(
+        label="Download ZIP"
+    )
 
     run_btn.click(
         fn=process_batch,
         inputs=files_in,
-        outputs=[gallery_out, zip_out],
+        outputs=[
+            gallery_out,
+            zip_out
+        ]
     )
 
+
+# =========================
+# START
+# =========================
 if __name__ == "__main__":
-    demo.launch()
+
+    demo.launch(
+        server_name="0.0.0.0",
+        server_port=7860
+    )