Update app.py

app.py

@@ -3,10 +3,12 @@ import subprocess
 import os
 import shutil
 from pathlib import Path
-from inference_coz_single import recursive_multiscale_sr
-from PIL import Image, ImageDraw
 import spaces
 
+# import the updated recursive_multiscale_sr that expects a list of centers
+from inference_coz_single import recursive_multiscale_sr
+
+from PIL import Image, ImageDraw
 
 # ------------------------------------------------------------------
 # CONFIGURE THESE PATHS TO MATCH YOUR PROJECT STRUCTURE
@@ -15,6 +17,7 @@ import spaces
 INPUT_DIR   = "samples"
 OUTPUT_DIR  = "inference_results/coz_vlmprompt"
 
+
 # ------------------------------------------------------------------
 # HELPER: Resize & center-crop to 512, preserving aspect ratio
 # ------------------------------------------------------------------
@@ -35,69 +38,141 @@ def resize_and_center_crop(img: Image.Image, size: int) -> Image.Image:
 
 
 # ------------------------------------------------------------------
-# HELPER: Draw four concentric, centered rectangles on a 512×512 image
+# HELPER: Draw four true “nested” rectangles, matching the SR logic
 # ------------------------------------------------------------------
 
-def make_preview_with_boxes(image_path: str, scale_option: str) -> Image.Image:
+def make_preview_with_boxes(
+    image_path: str,
+    scale_option: str,
+    cx_norm: float,
+    cy_norm: float,
+) -> Image.Image:
     """
-    1) Open the uploaded image from disk.
-    2) Resize & center-crop it to exactly 512×512.
-    3) Depending on scale_option ("1x","2x","4x"), compute four rectangle sizes:
-       - "1x": [512, 512, 512, 512]
-       - "2x": [256, 128, 64, 32]
-       - "4x": [128, 64, 32, 16]
-    4) Draw each of those four rectangles (outline only), all centered.
-    5) Return the modified PIL image.
+    1) Open the uploaded image, resize & center-crop to 512×512.
+    2) Let scale_int = int(scale_option.replace("x","")).
+       Then the four nested crop‐sizes (in pixels) are:
+         size[0] = 512 / (scale_int^1),
+         size[1] = 512 / (scale_int^2),
+         size[2] = 512 / (scale_int^3),
+         size[3] = 512 / (scale_int^4).
+    3) Iteratively compute each crop’s top-left in “original 512×512” space:
+       - Start with prev_tl = (0,0), prev_size = 512.
+       - For i in [0..3]:
+           center_abs_x = prev_tl_x + cx_norm * prev_size
+           center_abs_y = prev_tl_y + cy_norm * prev_size
+           unc_x0 = center_abs_x - (size[i]/2)
+           unc_y0 = center_abs_y - (size[i]/2)
+           clamp x0 ∈ [prev_tl_x, prev_tl_x + prev_size - size[i]]
+                 y0 ∈ [prev_tl_y, prev_tl_y + prev_size - size[i]]
+           Draw a rectangle from (x0, y0) to (x0 + size[i], y0 + size[i]).
+           Then set prev_tl = (x0, y0), prev_size = size[i].
+    4) Return the PIL image with those four truly nested outlines.
     """
     try:
         orig = Image.open(image_path).convert("RGB")
     except Exception as e:
-        # If something fails, return a plain 512×512 gray image as fallback
+        # On error, return a gray 512×512 with the error text
         fallback = Image.new("RGB", (512, 512), (200, 200, 200))
         draw = ImageDraw.Draw(fallback)
         draw.text((20, 20), f"Error:\n{e}", fill="red")
         return fallback
 
-    # 1. Resize & center-crop to 512×512
-    base = resize_and_center_crop(orig, 512)  # now `base.size == (512,512)`
+    # 1) Resize & center-crop to 512×512
+    base = resize_and_center_crop(orig, 512)
 
-    # 2. Determine the four box sizes
-    scale_int = int(scale_option.replace("x", ""))  # e.g. "2x" -> 2
-    if scale_int == 1:
+    # 2) Compute the four nested crop‐sizes
+    scale_int = int(scale_option.replace("x", ""))  # e.g. "4x" → 4
+    if scale_int <= 1:
+        # If 1×, then all “nested” sizes are 512 (no real nesting)
         sizes = [512, 512, 512, 512]
     else:
-        # For scale=2: sizes = [512//2, 512//(2*2), 512//(2*4), 512//(2*8)] -> [256,128,64,32]
-        # For scale=4: sizes = [512//4, 512//(4*2), 512//(4*4), 512//(4*8)] -> [128,64,32,16]
-        sizes = [512 // (scale_int * (2 ** i)) for i in range(4)]
+        sizes = [
+            512 // (scale_int ** (i + 1))
+            for i in range(4)
+        ]
+        # e.g. if scale_int=4 → sizes = [128, 32, 8, 2]
 
     draw = ImageDraw.Draw(base)
-
-    # 3. Outline color cycle (you can change these or use just one color)
     colors = ["red", "lime", "cyan", "yellow"]
-    width = 3  # thickness of each rectangle’s outline
+    width = 3
+
+    # 3) Iteratively compute nested rectangles
+    prev_tl_x, prev_tl_y = 0.0, 0.0
+    prev_size = 512.0
+
+    for idx, crop_size in enumerate(sizes):
+        # 3.a) Where is the “normalized center” in this current 512×512 region?
+        center_abs_x = prev_tl_x + (cx_norm * prev_size)
+        center_abs_y = prev_tl_y + (cy_norm * prev_size)
+
+        # 3.b) Unclamped top-left for this crop
+        unc_x0 = center_abs_x - (crop_size / 2.0)
+        unc_y0 = center_abs_y - (crop_size / 2.0)
+
+        # 3.c) Clamp so the crop window stays inside [prev_tl .. prev_tl + prev_size]
+        min_x0 = prev_tl_x
+        max_x0 = prev_tl_x + prev_size - crop_size
+        min_y0 = prev_tl_y
+        max_y0 = prev_tl_y + prev_size - crop_size
+
+        x0 = max(min_x0, min(unc_x0, max_x0))
+        y0 = max(min_y0, min(unc_y0, max_y0))
+        x1 = x0 + crop_size
+        y1 = y0 + crop_size
+
+        # Draw the rectangle (cast to int for pixels)
+        draw.rectangle(
+            [(int(x0), int(y0)), (int(x1), int(y1))],
+            outline=colors[idx % len(colors)],
+            width=width
+        )
 
-    for idx, s in enumerate(sizes):
-        # Compute top-left corner so that box is centered in 512×512
-        x0 = (512 - s) // 2
-        y0 = (512 - s) // 2
-        x1 = x0 + s
-        y1 = y0 + s
-        draw.rectangle([(x0, y0), (x1, y1)], outline=colors[idx % len(colors)], width=width)
+        # 3.d) Update for the next iteration
+        prev_tl_x, prev_tl_y = x0, y0
+        prev_size = crop_size
 
     return base
 
 
+# ------------------------------------------------------------------
+# HELPER FUNCTION FOR INFERENCE (build a list of identical centers)
+# ------------------------------------------------------------------
+
 @spaces.GPU(duration=120)
-def run_with_upload(uploaded_image_path, upscale_option):
+def run_with_upload(
+    uploaded_image_path: str,
+    upscale_option: str,
+    cx_norm: float,
+    cy_norm: float,
+):
+    """
+    - upscale_option: "1x" / "2x" / "4x"
+    - cx_norm, cy_norm: normalized center coordinates in [0,1]
+    The underlying `recursive_multiscale_sr` expects a list of centers
+    of length rec_num (default 4). We replicate (cx_norm, cy_norm) four times.
+    """
+    if uploaded_image_path is None:
+        return []
 
-  upscale_value = upscale_option.replace("x", "")  # e.g. "2x" → "2"
-  
-  return recursive_multiscale_sr(uploaded_image_path, int(upscale_value))[0]
+    upscale_value = int(upscale_option.replace("x", ""))
+    rec_num = 4  # match the SR pipeline’s default recursion depth
 
+    centers = [(cx_norm, cy_norm)] * rec_num
+
+    # Call the modified SR function
+    sr_list, _ = recursive_multiscale_sr(
+        uploaded_image_path,
+        upscale=upscale_value,
+        rec_num=rec_num,
+        centers=centers,
+    )
+
+    # Return the list of PIL images (Gradio Gallery expects a list)
+    return sr_list
 
 
 # ------------------------------------------------------------------
-# BUILD THE GRADIO INTERFACE (with updated callbacks)
+# BUILD THE GRADIO INTERFACE (two sliders + correct preview)
 # ------------------------------------------------------------------
 
 css = """
@@ -141,19 +216,35 @@ with gr.Blocks(css=css) as demo:
                     show_label=False
                 )
 
-                # 3) Button to launch inference
+                # 3) Two sliders for normalized center (0..1)
+                center_x = gr.Slider(
+                    label="Center X (normalized)",
+                    minimum=0.0,
+                    maximum=1.0,
+                    step=0.01,
+                    value=0.5
+                )
+                center_y = gr.Slider(
+                    label="Center Y (normalized)",
+                    minimum=0.0,
+                    maximum=1.0,
+                    step=0.01,
+                    value=0.5
+                )
+
+                # 4) Button to launch inference
                 run_button = gr.Button("Chain-of-Zoom it")
 
-                # 4) Show the 512×512 preview with four centered rectangles
+                # 5) Preview (512×512 + four truly nested boxes)
                 preview_with_box = gr.Image(
-                    label="Preview (512×512 with centered boxes)",
-                    type="pil",        # we’ll return a PIL.Image from our function
+                    label="Preview (512×512 with nested boxes)",
+                    type="pil",
                     interactive=False
                 )
 
 
             with gr.Column():
-                # 5) Gallery to display multiple output images
+                # 6) Gallery to display multiple output images
                 output_gallery = gr.Gallery(
                     label="Inference Results",
                     show_label=True,
@@ -162,39 +253,51 @@ with gr.Blocks(css=css) as demo:
                 )
 
         # ------------------------------------------------------------------
-        # CALLBACK #1: Whenever the user uploads or changes the radio, update preview
+        # CALLBACK #1: update the preview whenever inputs change
         # ------------------------------------------------------------------
 
-        def update_preview(img_path, scale_opt):
+        def update_preview(
+            img_path: str,
+            scale_opt: str,
+            cx: float,
+            cy: float
+        ) -> Image.Image | None:
             """
-            If there's no image uploaded yet, return None (Gradio will show blank).
-            Otherwise, draw the resized 512×512 + four boxes and return it.
+            If no image uploaded, show blank; otherwise, draw four nested boxes
+            exactly as the SR pipeline would crop at each recursion.
             """
             if img_path is None:
                 return None
-            return make_preview_with_boxes(img_path, scale_opt)
+            return make_preview_with_boxes(img_path, scale_opt, cx, cy)
 
-        # When the user uploads a new file:
         upload_image.change(
             fn=update_preview,
-            inputs=[upload_image, upscale_radio],
+            inputs=[upload_image, upscale_radio, center_x, center_y],
             outputs=[preview_with_box]
         )
-
-        # Also trigger preview redraw if they switch 1×/2×/4× after uploading:
         upscale_radio.change(
             fn=update_preview,
-            inputs=[upload_image, upscale_radio],
+            inputs=[upload_image, upscale_radio, center_x, center_y],
+            outputs=[preview_with_box]
+        )
+        center_x.change(
+            fn=update_preview,
+            inputs=[upload_image, upscale_radio, center_x, center_y],
+            outputs=[preview_with_box]
+        )
+        center_y.change(
+            fn=update_preview,
+            inputs=[upload_image, upscale_radio, center_x, center_y],
             outputs=[preview_with_box]
         )
 
         # ------------------------------------------------------------------
-        # CALLBACK #2: When “Chain-of-Zoom it” is clicked, run inference
+        # CALLBACK #2: on button‐click, run the SR pipeline
         # ------------------------------------------------------------------
 
         run_button.click(
             fn=run_with_upload,
-            inputs=[upload_image, upscale_radio],
+            inputs=[upload_image, upscale_radio, center_x, center_y],
             outputs=[output_gallery]
         )
 
@@ -203,5 +306,4 @@ with gr.Blocks(css=css) as demo:
 # START THE GRADIO SERVER
 # ------------------------------------------------------------------
 
-# 🔧 2) launch as usual
-demo.launch(share=True)
+demo.launch(share=True, debug=True)