import gradio as gr
import subprocess
import os
import shutil
from pathlib import Path
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
# ------------------------------------------------------------------
INPUT_DIR = "samples"
OUTPUT_DIR = "inference_results/coz_vlmprompt"
# ------------------------------------------------------------------
# HELPER: Resize & center-crop to 512, preserving aspect ratio
# ------------------------------------------------------------------
def resize_and_center_crop(img: Image.Image, size: int) -> Image.Image:
"""
Resize the input PIL image so that its shorter side == `size`,
then center-crop to exactly (size x size).
"""
w, h = img.size
scale = size / min(w, h)
new_w, new_h = int(w * scale), int(h * scale)
img = img.resize((new_w, new_h), Image.LANCZOS)
left = (new_w - size) // 2
top = (new_h - size) // 2
return img.crop((left, top, left + size, top + size))
# ------------------------------------------------------------------
# HELPER: Draw four true “nested” rectangles, matching the SR logic
# ------------------------------------------------------------------
def make_preview_with_boxes(
image_path: str,
scale_option: str,
cx_norm: float,
cy_norm: float,
) -> tuple[Image.Image, list[tuple[float, float]]]:
"""
Returns:
- The preview image with drawn boxes.
- A list of (cx_norm, cy_norm) for each box (normalized to 512×512).
"""
try:
orig = Image.open(image_path).convert("RGB")
except Exception as e:
fallback = Image.new("RGB", (512, 512), (200, 200, 200))
ImageDraw.Draw(fallback).text((20, 20), f"Error:\n{e}", fill="red")
return fallback, []
base = resize_and_center_crop(orig, 512)
scale_int = int(scale_option.replace("x", ""))
if scale_int <= 1:
sizes = [512.0, 512.0, 512.0, 512.0]
else:
sizes = [512.0 / (scale_int ** (i + 1)) for i in range(4)]
draw = ImageDraw.Draw(base)
colors = ["red", "lime", "cyan", "yellow"]
width = 3
abs_cx = cx_norm * 512.0
abs_cy = cy_norm * 512.0
prev_x0, prev_y0, prev_size = 0.0, 0.0, 512.0
centers: list[tuple[float, float]] = []
for i, crop_size in enumerate(sizes):
x0 = abs_cx - (crop_size / 2.0)
y0 = abs_cy - (crop_size / 2.0)
min_x0 = prev_x0
max_x0 = prev_x0 + prev_size - crop_size
min_y0 = prev_y0
max_y0 = prev_y0 + prev_size - crop_size
x0 = max(min_x0, min(x0, max_x0))
y0 = max(min_y0, min(y0, max_y0))
x1 = x0 + crop_size
y1 = y0 + crop_size
draw.rectangle([(int(round(x0)), int(round(y0))),
(int(round(x1)), int(round(y1)))],
outline=colors[i % len(colors)], width=width)
# --- compute normalized center of this box ---
cx_box = ((x0 - prev_x0) + crop_size / 2.0) / float(prev_size)
cy_box = ((y0 - prev_y0) + crop_size / 2.0) / float(prev_size)
centers.append((cx_box, cy_box))
prev_x0, prev_y0, prev_size = x0, y0, crop_size
return base, centers
# ------------------------------------------------------------------
# HELPER FUNCTION FOR INFERENCE (build a list of identical centers)
# ------------------------------------------------------------------
@spaces.GPU()
def run_with_upload(
uploaded_image_path: str,
upscale_option: str,
cx_norm: float,
cy_norm: float,
):
"""
Perform chain-of-zoom super-resolution on a given image, using recursive multi-scale upscaling centered on a specific point.
This function enhances a given image by progressively zooming into a specific point, using a recursive deep super-resolution model.
Args:
uploaded_image_path (str): Path to the input image file on disk.
upscale_option (str): The desired upscale factor as a string. Valid options are "1x", "2x", and "4x".
- "1x" means no upscaling.
- "2x" means 2× enlargement per zoom step.
- "4x" means 4× enlargement per zoom step.
cx_norm (float): Normalized X-coordinate (0 to 1) of the zoom center.
cy_norm (float): Normalized Y-coordinate (0 to 1) of the zoom center.
Returns:
list[PIL.Image.Image]: A list of progressively zoomed-in and super-resolved images at each recursion step (typically 4),
centered around the user-specified point.
Note:
The center point is repeated for each recursion level to maintain consistency during zooming.
This function uses a modified version of the `recursive_multiscale_sr` pipeline for inference.
"""
if uploaded_image_path is None:
return []
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
@spaces.GPU()
def magnify(
uploaded_image_path: str,
upscale_option: str,
centres: list
):
"""
Perform chain-of-zoom super-resolution on a given image, using recursive multi-scale upscaling centered on a specific point.
This function enhances a given image by progressively zooming into a specific point, using a recursive deep super-resolution model.
Args:
uploaded_image_path (str): Path to the input image file on disk.
upscale_option (str): The desired upscale factor as a string. Valid options are "1x", "2x", and "4x".
- "1x" means no upscaling.
- "2x" means 2× enlargement per zoom step.
- "4x" means 4× enlargement per zoom step.
centres (list): Normalized list of X-coordinate, Y-coordinate (0 to 1) of the zoom center.
Returns:
list[PIL.Image.Image]: A list of progressively zoomed-in and super-resolved images at each recursion step (typically 4),
centered around the user-specified point.
Note:
The center point is repeated for each recursion level to maintain consistency during zooming.
This function uses a modified version of the `recursive_multiscale_sr` pipeline for inference.
"""
if uploaded_image_path is None:
return []
upscale_value = int(upscale_option.replace("x", ""))
rec_num = len(centres)
# Call the modified SR function
sr_list, _ = recursive_multiscale_sr(
uploaded_image_path,
upscale=upscale_value,
rec_num=rec_num,
centers=centres,
)
# Return the list of PIL images (Gradio Gallery expects a list)
return sr_list
# ------------------------------------------------------------------
# BUILD THE GRADIO INTERFACE (two sliders + correct preview)
# ------------------------------------------------------------------
css = """
#col-container {
margin: 0 auto;
max-width: 1024px;
}
"""
with gr.Blocks(css=css) as demo:
session_centres = gr.State()
with gr.Column(elem_id="col-container"):
gr.HTML(
"""
Chain-of-Zoom – Extreme Super-Resolution via Scale Autoregression and Preference Alignment
