Prefer Paint workspace for live drawing with ImageEditor fallback

app.py

@@ -20,8 +20,10 @@ from collections import deque
 from io import StringIO, BytesIO
 
 HAS_IMAGE_EDITOR = hasattr(gr, "ImageEditor")
+HAS_PAINT = hasattr(gr, "Paint")
 HAS_BRUSH = hasattr(gr, "Brush")
 HAS_ERASER = hasattr(gr, "Eraser")
+HAS_REGION_WORKSPACE = HAS_PAINT or HAS_IMAGE_EDITOR
 
 # Model options — swap MODEL_NAME to reduce VRAM usage on GPUs with <= 8GB
 #
@@ -891,7 +893,7 @@ def _extract_editor_background(editor_value):
 
 def _to_rgba_image(obj):
     if isinstance(obj, dict):
-        for k in ("image", "layer", "composite", "background"):
+        for k in ("image", "layer", "composite", "background", "mask"):
             if k in obj:
                 return _to_rgba_image(obj[k])
         return None
@@ -909,6 +911,32 @@ def _to_rgba_image(obj):
         return Image.fromarray(arr.astype(np.uint8), mode="RGBA")
     return None
 
+def _to_mask_array(obj):
+    if obj is None:
+        return None
+    if isinstance(obj, dict):
+        for k in ("mask", "image", "layer", "composite", "background"):
+            if k in obj:
+                arr = _to_mask_array(obj[k])
+                if arr is not None:
+                    return arr
+        return None
+
+    if isinstance(obj, Image.Image):
+        arr = np.asarray(obj)
+    elif isinstance(obj, np.ndarray):
+        arr = obj
+    else:
+        return None
+
+    if arr.ndim == 2:
+        return arr > 0
+    if arr.ndim == 3:
+        if arr.shape[2] >= 4:
+            return arr[:, :, 3] > 0
+        return np.max(arr[:, :, :3], axis=2) > 0
+    return None
+
 def _locate_patch_bbox(base_image: Image.Image, patch_image: Image.Image):
     """Approximate patch location in base image using downscaled SSD search."""
     if base_image is None or patch_image is None:
@@ -1034,12 +1062,23 @@ def _editor_background_and_mask(editor_value):
     if not isinstance(editor_value, dict):
         return None, None
     background = _to_rgba_image(editor_value.get("background"))
+    if background is None:
+        background = _to_rgba_image(editor_value.get("image"))
     composite = _to_rgba_image(editor_value.get("composite"))
     layers = editor_value.get("layers") or []
     if background is None:
         if composite is None:
             return None, None
         background = composite
+
+    mask = _to_mask_array(editor_value.get("mask"))
+    if mask is not None:
+        if mask.shape[:2] != (background.height, background.width):
+            mask_img = Image.fromarray(mask.astype(np.uint8) * 255, mode="L")
+            nearest = Image.Resampling.NEAREST if hasattr(Image, "Resampling") else Image.NEAREST
+            mask = np.asarray(mask_img.resize((background.width, background.height), nearest)) > 0
+        return background, mask
+
     if not isinstance(layers, list) or not layers:
         return background, None
 
@@ -1096,11 +1135,17 @@ def _extract_new_drawn_regions(editor_value, base_size=None, base_image=None, co
     if background is None:
         return [], consumed_mask
 
-    # If there are no drawn layers, treat as explicit crop mode.
-    if not isinstance(layers, list) or not layers or mask is None:
+    has_layer_data = isinstance(layers, list) and len(layers) > 0
+    has_draw_data = (mask is not None) or has_layer_data
+
+    # If there are no draw layers/mask, treat as explicit crop mode.
+    if not has_draw_data:
         regions = _extract_selected_regions(editor_value, base_size=base_size, base_image=base_image)
         return regions, consumed_mask
 
+    if mask is None:
+        return [], consumed_mask
+
     if consumed_mask is None or not isinstance(consumed_mask, np.ndarray) or consumed_mask.shape != mask.shape:
         delta_mask = mask
     else:
@@ -1349,13 +1394,14 @@ with gr.Blocks(title="DeepSeek-OCR-2") as demo:
         with gr.Column(scale=3):
             workspace_hint = gr.Markdown("**Entire Page mode:** No drawing needed. Click **Extract** to process the full page.")
             gr.Markdown("**Image Workspace (full page + region selection)**")
-            if HAS_IMAGE_EDITOR:
+            if HAS_REGION_WORKSPACE:
                 editor_kwargs = {}
                 if HAS_BRUSH:
                     try:
+                        highlight = ("#2563eb", 0.35)
                         editor_kwargs["brush"] = gr.Brush(
-                            colors=["#2563eb"],
-                            default_color="#2563eb",
+                            colors=[highlight],
+                            default_color=highlight,
                             color_mode="fixed",
                             default_size=22,
                         )
@@ -1374,15 +1420,39 @@ with gr.Blocks(title="DeepSeek-OCR-2") as demo:
                         editor_kwargs["eraser"] = gr.Eraser(default_size=26)
                     except TypeError:
                         editor_kwargs["eraser"] = gr.Eraser()
-                region_editor = gr.ImageEditor(
-                    label="Image Workspace",
-                    show_label=False,
-                    type="pil",
-                    height=640,
-                    **editor_kwargs,
-                )
+                if HAS_PAINT:
+                    try:
+                        region_editor = gr.Paint(
+                            label="Image Workspace",
+                            show_label=False,
+                            type="pil",
+                            height=640,
+                            **editor_kwargs,
+                        )
+                    except TypeError:
+                        try:
+                            region_editor = gr.Paint(
+                                label="Image Workspace",
+                                show_label=False,
+                                height=640,
+                                **editor_kwargs,
+                            )
+                        except TypeError:
+                            region_editor = gr.Paint(
+                                label="Image Workspace",
+                                show_label=False,
+                                height=640,
+                            )
+                else:
+                    region_editor = gr.ImageEditor(
+                        label="Image Workspace",
+                        show_label=False,
+                        type="pil",
+                        height=640,
+                        **editor_kwargs,
+                    )
             else:
-                gr.Markdown("Region drawing requires a newer Gradio version with `ImageEditor` support.")
+                gr.Markdown("Region drawing requires a newer Gradio version with `Paint` or `ImageEditor` support.")
                 region_editor = gr.State(None)
 
         with gr.Column(scale=1):
@@ -1433,7 +1503,7 @@ with gr.Blocks(title="DeepSeek-OCR-2") as demo:
                 "examples/2022-0922 Section 14 Notes.png",
                 "examples/2022-0922 Section 15 Notes.png",
             ]
-            if HAS_IMAGE_EDITOR and region_editor is not None:
+            if HAS_REGION_WORKSPACE and region_editor is not None:
                 gr.Examples(
                     label="Image Examples (click thumbnail to load into workspace)",
                     examples=image_examples,
@@ -1504,7 +1574,7 @@ with gr.Blocks(title="DeepSeek-OCR-2") as demo:
     task.change(toggle_prompt, [task], [prompt])
     task.change(select_boxes, [task], [tabs])
     input_scope.change(toggle_scope_ui, [input_scope], [workspace_hint, selection_controls, selection_status, selected_regions_gallery])
-    if HAS_IMAGE_EDITOR and region_editor is not None:
+    if HAS_REGION_WORKSPACE and region_editor is not None:
         file_in.change(load_image_with_size, [file_in, page_selector], [region_editor, workspace_base_size, workspace_base_image])
         page_selector.change(load_image_with_size, [file_in, page_selector], [region_editor, workspace_base_size, workspace_base_image])
         region_editor.change(sync_workspace_state, [region_editor, workspace_base_image], [workspace_base_size, workspace_base_image])