Update app.py

app.py

@@ -136,11 +136,11 @@ enhancer = RealESRGANer(
     device=restoration_device,
 )
 
-# 2. GFPGAN Weights Setup
-gfpgan_path = f"{weights_dir}/GFPGANv1.3.pth"
+# 2. GFPGAN Weights Setup — version 1.4
+gfpgan_path = f"{weights_dir}/GFPGANv1.4.pth"
 ensure_weight_file(
     gfpgan_path,
-    "https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.3.pth",
+    "https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.4.pth",
 )
 
 gfpgan_restorer = GFPGANer(
@@ -227,6 +227,48 @@ def cap_image(image, max_side=1024):
     return image.resize((int(image.width * ratio), int(image.height * ratio)), Image.LANCZOS)
 
 
+def cleanup_vintage_image(image, scratch_mask=None, denoise_strength=0):
+    """CPU-only noise cleanup and optional user-mask-guided scratch repair.
+
+    Paint/prepare the mask with white over defects to repair and black elsewhere.
+    OpenCV inpainting works best on thin scratches and small blemishes, not large
+    missing areas. A denoise value of zero disables denoising.
+    """
+    if image is None:
+        raise gr.Error("Please upload an image first.")
+
+    image = cap_image(image, 1200)
+    rgb = np.asarray(image.convert("RGB"), dtype=np.uint8)
+
+    if scratch_mask is not None:
+        mask = np.asarray(
+            scratch_mask.convert("L").resize(image.size, Image.NEAREST),
+            dtype=np.uint8,
+        )
+        _, mask = cv2.threshold(mask, 127, 255, cv2.THRESH_BINARY)
+        if np.any(mask):
+            # Expand slightly to cover the edge of a painted scratch mask.
+            mask = cv2.dilate(mask, np.ones((3, 3), dtype=np.uint8), iterations=1)
+            bgr = cv2.cvtColor(rgb, cv2.COLOR_RGB2BGR)
+            bgr = cv2.inpaint(bgr, mask, 3, cv2.INPAINT_TELEA)
+            rgb = cv2.cvtColor(bgr, cv2.COLOR_BGR2RGB)
+
+    denoise_strength = int(denoise_strength)
+    if not 0 <= denoise_strength <= 10:
+        raise gr.Error("Denoise strength must be between 0 and 10.")
+    if denoise_strength > 0:
+        rgb = cv2.fastNlMeansDenoisingColored(
+            rgb,
+            None,
+            float(denoise_strength),
+            float(denoise_strength),
+            7,
+            21,
+        )
+
+    return Image.fromarray(rgb)
+
+
 def _colorfulness(image):
     rgb = np.array(image.convert("RGB"))
     R, G, B = rgb[:, :, 0], rgb[:, :, 1], rgb[:, :, 2]
@@ -258,17 +300,34 @@ def enhance_image(image, scale_factor, mode="Pure Super-Resolution (Safe for Gro
 
     try:
         if "Pure Super-Resolution" in mode:
-            # Pure Real-ESRGAN Upscaling & Deblurring (preserves exact facial features and group structures)
+            # Real-ESRGAN upscaling only; this mode does not run GFPGAN face restoration.
             output_bgr, _ = enhancer.enhance(input_bgr, outscale=target_scale)
         else:
-            # GFPGAN Face Deblurring + Real-ESRGAN Canvas Upscaling
+            # GFPGAN restores detected faces; Real-ESRGAN upsamples the canvas.
+            # In GFPGAN's clean architecture the `weight` kwarg is not used for
+            # blending, so blend explicitly here for consistent UI behavior:
+            # 0.0 = full restored result; larger values preserve more of the
+            # upscaled input and reduce the visible restoration effect.
             gfpgan_restorer.upscale = target_scale
-            _, _, output_bgr = gfpgan_restorer.enhance(
+            _, _, restored_bgr = gfpgan_restorer.enhance(
                 input_bgr,
                 has_aligned=False,
                 only_center_face=False,
                 paste_back=True,
-                weight=float(fidelity_weight)
+                weight=0.0,
+            )
+            fidelity = float(np.clip(fidelity_weight, 0.0, 0.7))
+            input_upscaled = cv2.resize(
+                input_bgr,
+                (restored_bgr.shape[1], restored_bgr.shape[0]),
+                interpolation=cv2.INTER_LANCZOS4,
+            )
+            output_bgr = cv2.addWeighted(
+                restored_bgr,
+                1.0 - fidelity,
+                input_upscaled,
+                fidelity,
+                0.0,
             )
 
         output_rgb = cv2.cvtColor(output_bgr, cv2.COLOR_BGR2RGB)
@@ -600,11 +659,37 @@ with gr.Blocks(
 
         # Enhance Tab UI
         with gr.Tab("✨ Enhance", id="tab_enhance"):
-            gr.Markdown("## Restore old blurry photos, sharpen features, and upscale up to 4x.")
+            gr.Markdown(
+                "## Restore faces, clean noise/small scratches, and upscale up to 4x.\n"
+                "The cleanup tool is CPU/OpenCV-based. Upload a grayscale mask with "
+                "white over scratches and black elsewhere."
+            )
             with gr.Row(equal_height=True):
                 with gr.Column(scale=1):
                     enh_input = gr.Image(label="Upload Image", type="pil", height=350)
-                    
+
+                    with gr.Accordion("Vintage Cleanup (CPU / OpenCV)", open=False):
+                        scratch_mask_input = gr.Image(
+                            label="Optional scratch mask (white = repair, black = keep)",
+                            type="pil",
+                            height=220,
+                        )
+                        denoise_strength = gr.Slider(
+                            minimum=0,
+                            maximum=10,
+                            value=0,
+                            step=1,
+                            label="Noise reduction (0 = off; higher can soften detail)",
+                        )
+                        cleanup_btn = gr.Button("Clean / Repair Image", variant="secondary")
+                        cleanup_preview = gr.Image(
+                            label="Cleaned Preview (use below to continue)",
+                            type="pil",
+                            interactive=False,
+                            height=300,
+                        )
+                        use_cleaned_btn = gr.Button("Use Cleaned Image for Enhancement")
+
                     enh_mode = gr.Radio(
                         choices=[
                             "Pure Super-Resolution (Safe for Group Photos)", 
@@ -616,8 +701,11 @@ with gr.Blocks(
                     
                     enh_scale = gr.Radio(choices=["2x", "4x"], value="4x", label="Upscale Factor")
                     enh_weight = gr.Slider(
-                        minimum=0.0, maximum=0.5, value=0.15, step=0.05, 
-                        label="Face Restoration Weight (Only applies in Refinement mode; keep low <= 0.2)"
+                        minimum=0.0,
+                        maximum=0.7,
+                        value=0.15,
+                        step=0.05,
+                        label="Original preservation (0 = strongest face restoration; higher = more original)",
                     )
                     enh_metrics = gr.Checkbox(label="📊 Compute quality metrics", value=False)
                     enh_btn = gr.Button("Enhance Image", variant="primary", size="lg")
@@ -625,10 +713,27 @@ with gr.Blocks(
                     enh_output = gr.Image(label="Result", type="pil", interactive=False, height=350)
                     enh_info = gr.Textbox(label="Size Info + Metrics", interactive=False, lines=4)
             
+            cleanup_btn.click(
+                fn=cleanup_vintage_image,
+                inputs=[enh_input, scratch_mask_input, denoise_strength],
+                outputs=[cleanup_preview],
+            )
+
+            def use_cleaned_image(image):
+                if image is None:
+                    raise gr.Error("Run cleanup first, or upload an image.")
+                return image
+
+            use_cleaned_btn.click(
+                fn=use_cleaned_image,
+                inputs=[cleanup_preview],
+                outputs=[enh_input],
+            )
+
             enh_btn.click(
-                fn=enhance_image, 
-                inputs=[enh_input, enh_scale, enh_mode, enh_weight, enh_metrics], 
-                outputs=[enh_output, enh_info]
+                fn=enhance_image,
+                inputs=[enh_input, enh_scale, enh_mode, enh_weight, enh_metrics],
+                outputs=[enh_output, enh_info],
             )
 
     # Inter-tab Event Routing Connections
@@ -636,4 +741,4 @@ with gr.Blocks(
     btn_col_to_enhance.click(fn=send_to_enhance, inputs=[col_output], outputs=[enh_input, tabs])
     btn_out_to_enhance.click(fn=send_to_enhance, inputs=[out_output], outputs=[enh_input, tabs])
 
-demo.queue().launch(show_error=True)
+demo.queue().launch(show_error=True)