Detecting and Editing Visual Objects with Gemini

🛠️ Option 1 – Gemini API via Vertex AI

Requirements:

Gen AI SDK environment variables:

• GOOGLE_GENAI_USE_VERTEXAI="True"
• GOOGLE_CLOUD_PROJECT="<PROJECT_ID>"
• GOOGLE_CLOUD_LOCATION="<LOCATION>"

💡 For preview models, the location must be set to global. For generally available models, we can choose the closest location among the Google model endpoint locations.

ℹ️ Learn more about setting up a project and a development environment.

🛠️ Option 2 – Gemini API via Google AI Studio

Requirement:

Gen AI SDK environment variables:

• GOOGLE_GENAI_USE_VERTEXAI="False"
• GOOGLE_API_KEY="<API_KEY>"

ℹ️ Learn more about getting a Gemini API key from Google AI Studio.

💡 You can store your environment configuration outside of the source code:

import os
import sys
from collections.abc import Callable

from google import genai

# Manual setup (leave unchanged if setup is environment-defined)

# @markdown **Which API: Vertex AI or Google AI Studio?**
GOOGLE_GENAI_USE_VERTEXAI = True  # @param {type: "boolean"}

# @markdown **Option A - Google Cloud project [+location]**
GOOGLE_CLOUD_PROJECT = ""  # @param {type: "string"}
GOOGLE_CLOUD_LOCATION = "global"  # @param {type: "string"}

# @markdown **Option B - Google AI Studio API key**
GOOGLE_API_KEY = ""  # @param {type: "string"}


def check_environment() -> bool:
    check_colab_user_authentication()
    return check_manual_setup() or check_vertex_ai() or check_colab() or check_local()


def check_manual_setup() -> bool:
    return check_define_env_vars(
        GOOGLE_GENAI_USE_VERTEXAI,
        GOOGLE_CLOUD_PROJECT.strip(),  # Might have been pasted with line return
        GOOGLE_CLOUD_LOCATION,
        GOOGLE_API_KEY,
    )


def check_vertex_ai() -> bool:
    # Workbench and Colab Enterprise
    match os.getenv("VERTEX_PRODUCT", ""):
        case "WORKBENCH_INSTANCE":
            pass
        case "COLAB_ENTERPRISE":
            if not running_in_colab_env():
                return False
        case _:
            return False

    return check_define_env_vars(
        True,
        os.getenv("GOOGLE_CLOUD_PROJECT", ""),
        os.getenv("GOOGLE_CLOUD_REGION", ""),
        "",
    )


def check_colab() -> bool:
    if not running_in_colab_env():
        return False

    # Colab Enterprise was checked before, so this is Colab only
    from google.colab import auth as colab_auth  # type: ignore

    colab_auth.authenticate_user()

    # Use Colab Secrets (🗝️ icon in left panel) to store the environment variables
    # Secrets are private, visible only to you and the notebooks that you select
    # - Vertex AI: Store your settings as secrets
    # - Google AI: Directly import your Gemini API key from the UI
    vertexai, project, location, api_key = get_vars(get_colab_secret)

    return check_define_env_vars(vertexai, project, location, api_key)


def check_local() -> bool:
    vertexai, project, location, api_key = get_vars(os.getenv)

    return check_define_env_vars(vertexai, project, location, api_key)


def running_in_colab_env() -> bool:
    # Colab or Colab Enterprise
    return "google.colab" in sys.modules


def check_colab_user_authentication() -> None:
    if running_in_colab_env():
        from google.colab import auth as colab_auth  # type: ignore

        colab_auth.authenticate_user()


def get_colab_secret(secret_name: str, default: str) -> str:
    from google.colab import errors, userdata  # type: ignore

    try:
        return userdata.get(secret_name)
    except errors.SecretNotFoundError:
        return default


def disable_colab_cell_scrollbar() -> None:
    if running_in_colab_env():
        from google.colab import output  # type: ignore

        output.no_vertical_scroll()


def get_vars(getenv: Callable[[str, str], str]) -> tuple[bool, str, str, str]:
    # Limit getenv calls to the minimum (may trigger UI confirmation for secret access)
    vertexai_str = getenv("GOOGLE_GENAI_USE_VERTEXAI", "")
    if vertexai_str:
        vertexai = vertexai_str.lower() in ["true", "1"]
    else:
        vertexai = bool(getenv("GOOGLE_CLOUD_PROJECT", ""))

    project = getenv("GOOGLE_CLOUD_PROJECT", "") if vertexai else ""
    location = getenv("GOOGLE_CLOUD_LOCATION", "") if project else ""
    api_key = getenv("GOOGLE_API_KEY", "") if not project else ""

    return vertexai, project, location, api_key


def check_define_env_vars(
    vertexai: bool,
    project: str,
    location: str,
    api_key: str,
) -> bool:
    match (vertexai, bool(project), bool(location), bool(api_key)):
        case (True, True, _, _):
            # Vertex AI - Google Cloud project [+location]
            location = location or "global"
            define_env_vars(vertexai, project, location, "")
        case (True, False, _, True):
            # Vertex AI - API key
            define_env_vars(vertexai, "", "", api_key)
        case (False, _, _, True):
            # Google AI Studio - API key
            define_env_vars(vertexai, "", "", api_key)
        case _:
            return False

    return True


def define_env_vars(vertexai: bool, project: str, location: str, api_key: str) -> None:
    os.environ["GOOGLE_GENAI_USE_VERTEXAI"] = str(vertexai)
    os.environ["GOOGLE_CLOUD_PROJECT"] = project
    os.environ["GOOGLE_CLOUD_LOCATION"] = location
    os.environ["GOOGLE_API_KEY"] = api_key


def check_configuration(client: genai.Client) -> None:
    service = "Vertex AI" if client.vertexai else "Google AI Studio"
    print(f"✅ Using the {service} API", end="")

    if client._api_client.project:
        print(f' with project "{client._api_client.project[:7]}…"', end="")
        print(f' in location "{client._api_client.location}"')
    elif client._api_client.api_key:
        api_key = client._api_client.api_key
        print(f' with API key "{api_key[:5]}…{api_key[-5:]}"', end="")
        print(f" (in case of error, make sure it was created for {service})")


print("✅ Environment functions defined")

from dataclasses import dataclass
from enum import StrEnum

Url = str


class Source(StrEnum):
    incunable = "https://tile.loc.gov/image-services/iiif/service:rbc:rbc0001:2014:2014rosen0487:0165/full/pct:25/0/default.jpg"
    engravings = "https://tile.loc.gov/image-services/iiif/service:gdc:gdcscd:00:34:07:66:92:1:00340766921:0121/full/pct:50/0/default.jpg"
    museum_guidebook = "https://tile.loc.gov/image-services/iiif/service:rbc:rbc0001:2014:2014gen34181:0033/full/pct:75/0/default.jpg"
    denver_illustrated = "https://tile.loc.gov/image-services/iiif/service:gdc:gdclccn:rc:01:00:04:94:rc01000494:0051/full/pct:50/0/default.jpg"
    physics_textbook = "https://tile.loc.gov/image-services/iiif/service:gdc:gdcscd:00:03:64:87:31:8:00036487318:0103/full/pct:50/0/default.jpg"
    portrait_miniatures = "https://tile.loc.gov/image-services/iiif/service:rbc:rbc0001:2024:2024rosen013592v02:0249/full/pct:50/0/default.jpg"
    wizard_of_oz_drawings = "https://tile.loc.gov/image-services/iiif/service:rbc:rbc0001:2006:2006gen32405:0048/full/pct:25/0/default.jpg"
    paintings = "https://images.unsplash.com/photo-1714146681164-f26fed839692?h=1440"
    alice_drawing = "https://images.unsplash.com/photo-1630595011903-689853b04ee2?h=800"
    book = "https://images.unsplash.com/photo-1643451533573-ee364ba6e330?h=800"
    manual = "https://images.unsplash.com/photo-1623666936367-a100f62ba9b7?h=800"
    electronics = "https://images.unsplash.com/photo-1757397584789-8b2c5bfcdbc3?h=1440"


@dataclass
class SourceMetadata:
    title: str
    webpage_url: Url
    credit_line: str


LOC = "Library of Congress"
LOC_RARE_BOOKS = "Library of Congress, Rare Book and Special Collections Division"
LOC_MEETING_FRONTIERS = "Library of Congress, Meeting of Frontiers"

metadata_by_source: dict[Source, SourceMetadata] = {
    Source.incunable: SourceMetadata(
        "Vergaderinge der historien van Troy (1485)",
        "https://www.loc.gov/resource/rbc0001.2014rosen0487/?sp=165",
        LOC_RARE_BOOKS,
    ),
    Source.engravings: SourceMetadata(
        "Harper's illustrated catalogue (1847)",
        "https://www.loc.gov/resource/gdcscd.00340766921/?sp=121",
        LOC,
    ),
    Source.museum_guidebook: SourceMetadata(
        "Barnum's American Museum illustrated (1850)",
        "https://www.loc.gov/resource/rbc0001.2014gen34181/?sp=33",
        LOC_RARE_BOOKS,
    ),
    Source.denver_illustrated: SourceMetadata(
        "Denver illustrated (1893)",
        "https://www.loc.gov/resource/gdclccn.rc01000494/?sp=51",
        LOC_MEETING_FRONTIERS,
    ),
    Source.physics_textbook: SourceMetadata(
        "Lessons in physics (1916)",
        "https://www.loc.gov/resource/gdcscd.00036487318/?sp=103",
        LOC,
    ),
    Source.portrait_miniatures: SourceMetadata(
        "The history of portrait miniatures (1904)",
        "https://www.loc.gov/resource/rbc0001.2024rosen013592v02/?sp=249",
        LOC_RARE_BOOKS,
    ),
    Source.wizard_of_oz_drawings: SourceMetadata(
        "The wonderful Wizard of Oz (1899)",
        "https://www.loc.gov/resource/rbc0001.2006gen32405/?sp=48",
        LOC_RARE_BOOKS,
    ),
    Source.paintings: SourceMetadata(
        "Open book showing paintings by Vincent van Gogh",
        "https://unsplash.com/photos/9hD7qrxICag",
        "Photo by Trung Manh cong on Unsplash",
    ),
    Source.alice_drawing: SourceMetadata(
        "Open book showing an illustration and text from Alice's Adventures in Wonderland",
        "https://unsplash.com/photos/bewzr_Q9u2o",
        "Photo by Brett Jordan on Unsplash",
    ),
    Source.book: SourceMetadata(
        "Open book showing two botanical illustrations",
        "https://unsplash.com/photos/4IDqcNj827I",
        "Photo by Ranurte on Unsplash",
    ),
    Source.manual: SourceMetadata(
        "Open user manual for vintage camera",
        "https://unsplash.com/photos/aaFU96eYASk",
        "Photo by Annie Spratt on Unsplash",
    ),
    Source.electronics: SourceMetadata(
        "Circuit board with electronic components",
        "https://unsplash.com/photos/Aqa1pHQ57pw",
        "Photo by Albert Stoynov on Unsplash",
    ),
}

print("✅ Test images defined")

from enum import auto
from pathlib import Path
from typing import Any, cast

import IPython.display
import matplotlib.pyplot as plt
import pydantic
import tenacity
from google.genai.errors import ClientError
from google.genai.types import (
    FinishReason,
    GenerateContentConfig,
    GenerateContentResponse,
    PIL_Image,
    ThinkingConfig,
    ThinkingLevel,
)


# Multimodal models with spatial understanding and structured outputs
class MultimodalModel(StrEnum):
    # Generally Available (GA)
    GEMINI_2_5_FLASH = "gemini-2.5-flash"
    GEMINI_2_5_PRO = "gemini-2.5-pro"
    # Preview
    GEMINI_3_FLASH_PREVIEW = "gemini-3-flash-preview"
    GEMINI_3_1_PRO_PREVIEW = "gemini-3.1-pro-preview"
    # Default model used for object detection
    DEFAULT = GEMINI_3_FLASH_PREVIEW


# Image generation and editing models
class ImageModel(StrEnum):
    # Generally Available (GA)
    GEMINI_2_5_FLASH_IMAGE = "gemini-2.5-flash-image"  # Nano Banana 🍌
    # Preview
    GEMINI_3_PRO_IMAGE_PREVIEW = "gemini-3-pro-image-preview"  # Nano Banana Pro 🍌
    # Default model used for image editing
    DEFAULT = GEMINI_2_5_FLASH_IMAGE


Model = MultimodalModel | ImageModel


def generate_content(
    contents: list[Any],
    model: Model,
    config: GenerateContentConfig | None,
    should_display_response_info: bool = False,
) -> GenerateContentResponse | None:
    response = None
    client = check_client_for_model(model)

    for attempt in get_retrier():
        with attempt:
            response = client.models.generate_content(
                model=model.value,
                contents=contents,
                config=config,
            )
    if should_display_response_info:
        display_response_info(response, config)

    return response


def check_client_for_model(model: Model) -> genai.Client:
    if (
        model.value.endswith("-preview")
        and client.vertexai
        and client._api_client.location != "global"
    ):
        # Preview models are only available on the "global" location
        return genai.Client(location="global")

    return client


def display_response_info(
    response: GenerateContentResponse | None,
    config: GenerateContentConfig | None,
) -> None:
    if response is None:
        print("❌ No response")
        return

    if usage_metadata := response.usage_metadata:
        if usage_metadata.prompt_token_count:
            print(f"Input tokens   : {usage_metadata.prompt_token_count:9,d}")
        if usage_metadata.candidates_token_count:
            print(f"Output tokens  : {usage_metadata.candidates_token_count:9,d}")
        if usage_metadata.thoughts_token_count:
            print(f"Thoughts tokens: {usage_metadata.thoughts_token_count:9,d}")

    if (
        config is not None
        and config.response_mime_type == "application/json"
        and response.parsed is None
    ):
        print("❌ Could not parse the JSON response")
        return
    if not response.candidates:
        print("❌ No `response.candidates`")
        return
    if (finish_reason := response.candidates[0].finish_reason) != FinishReason.STOP:
        print(f"❌ {finish_reason = }")
    if not response.text:
        print("❌ No `response.text`")
        return


def generate_image(
    sources: list[PIL_Image],
    prompt: str,
    model: ImageModel,
    config: GenerateContentConfig | None = None,
) -> PIL_Image | None:
    contents = [*sources, prompt.strip()]

    response = generate_content(contents, model, config)

    return check_get_output_image_from_response(response)


def check_get_output_image_from_response(
    response: GenerateContentResponse | None,
) -> PIL_Image | None:
    if response is None:
        print("❌ No `response`")
        return None
    if not response.candidates:
        print("❌ No `response.candidates`")
        if response.prompt_feedback:
            if block_reason := response.prompt_feedback.block_reason:
                print(f"{block_reason = :s}")
            if block_reason_message := response.prompt_feedback.block_reason_message:
                print(f"{block_reason_message = }")
        return None
    if not (content := response.candidates[0].content):
        print("❌ No `response.candidates[0].content`")
        return None
    if not (parts := content.parts):
        print("❌ No `response.candidates[0].content.parts`")
        return None

    output_image: PIL_Image | None = None
    for part in parts:
        if part.text:
            display_markdown(part.text)
            continue
        sdk_image = part.as_image()
        assert sdk_image is not None
        output_image = sdk_image._pil_image
        assert output_image is not None
        break  # There should be a single image

    return output_image


def get_thinking_config(model: Model) -> ThinkingConfig | None:
    match model:
        case MultimodalModel.GEMINI_2_5_FLASH:
            return ThinkingConfig(thinking_budget=0)
        case MultimodalModel.GEMINI_2_5_PRO:
            return ThinkingConfig(thinking_budget=128, include_thoughts=False)
        case MultimodalModel.GEMINI_3_FLASH_PREVIEW:
            return ThinkingConfig(thinking_level=ThinkingLevel.MINIMAL)
        case MultimodalModel.GEMINI_3_1_PRO_PREVIEW:
            return ThinkingConfig(thinking_level=ThinkingLevel.LOW)
        case _:
            return None  # Default


def display_markdown(markdown: str) -> None:
    IPython.display.display(IPython.display.Markdown(markdown))


def display_image(image: PIL_Image) -> None:
    IPython.display.display(image)


def get_retrier() -> tenacity.Retrying:
    return tenacity.Retrying(
        stop=tenacity.stop_after_attempt(7),
        wait=tenacity.wait_incrementing(start=10, increment=1),
        retry=should_retry_request,
        reraise=True,
    )


def should_retry_request(retry_state: tenacity.RetryCallState) -> bool:
    if not retry_state.outcome:
        return False
    err = retry_state.outcome.exception()
    if not isinstance(err, ClientError):
        return False
    print(f"❌ ClientError {err.code}: {err.message}")

    retry = False
    match err.code:
        case 400 if err.message is not None and " try again " in err.message:
            # Workshop: first time access to Cloud Storage (service agent provisioning)
            retry = True
        case 429:
            # Workshop: temporary project with 1 QPM quota
            retry = True
    print(f"🔄 Retry: {retry}")

    return retry


print("✅ Helpers defined")

import io
import urllib.request
from collections.abc import Iterator
from dataclasses import field
from datetime import datetime

import PIL.Image
from google.genai.types import Part, PartMediaResolutionLevel
from PIL.PngImagePlugin import PngInfo

OBJECT_DETECTION_PROMPT = """
Detect every illustration within the book photo and extract the following data for each:
- `box_2d`: Bounding box coordinates of the illustration only (ignoring any caption).
- `caption`: Verbatim caption or legend such as "Figure 1". Use "" if not found.
- `label`: Single-word label describing the illustration. Use "" if not found.
"""

# Margin added to detected/cropped objects, giving more context for a better understanding of spatial distortions
CROP_MARGIN_PX = 10

# Set to True to save each generated image
SAVE_GENERATED_IMAGES = False
OUTPUT_IMAGES_PATH = Path("./object_detection_and_editing")


# Matching class for structured output generation
class DetectedObject(pydantic.BaseModel):
    box_2d: list[int]
    caption: str
    label: str


# Misc data classes
InputImage = Path | Url
DetectedObjects = list[DetectedObject]
WorkflowStepImages = list[PIL_Image]


class WorkflowStep(StrEnum):
    SOURCE = auto()
    CROPPED = auto()
    RESTORED = auto()
    COLORIZED = auto()
    CINEMATIZED = auto()


@dataclass
class VisualObjectWorkflow:
    source_image: PIL_Image
    detected_objects: DetectedObjects
    images_by_step: dict[WorkflowStep, WorkflowStepImages] = field(default_factory=dict)

    def __post_init__(self) -> None:
        denormalize_bounding_boxes(self)


workflow_by_image: dict[InputImage, VisualObjectWorkflow] = {}


def denormalize_bounding_boxes(self: VisualObjectWorkflow) -> None:
    """Convert the box_2d coordinates.
    - Before: [y1, x1, y2, x2] normalized to 0-1000, as returned by Gemini
    - After:  [x1, y1, x2, y2] in source_image coordinates, as used in Pillow
    """

    def to_image_coord(coord: int, dim: int) -> int:
        return int(coord * dim / 1000 + 0.5)

    w, h = self.source_image.size
    for obj in self.detected_objects:
        y1, x1, y2, x2 = obj.box_2d
        x1, x2 = to_image_coord(x1, w), to_image_coord(x2, w)
        y1, y2 = to_image_coord(y1, h), to_image_coord(y2, h)
        obj.box_2d = [x1, y1, x2, y2]


def detect_objects(
    image: InputImage,
    prompt: str = OBJECT_DETECTION_PROMPT,
    model: MultimodalModel = MultimodalModel.DEFAULT,
    config: GenerateContentConfig | None = None,
    media_resolution: PartMediaResolutionLevel | None = None,
    display_results: bool = True,
) -> None:
    display_image_source_info(image)
    pil_image, content_part = get_pil_image_and_part(image, model, media_resolution)
    prompt = prompt.strip()
    contents = [content_part, prompt]
    config = config or get_object_detection_config(model)

    response = generate_content(contents, model, config)

    if response is not None and response.parsed is not None:
        detected_objects = cast(DetectedObjects, response.parsed)
    else:
        detected_objects = DetectedObjects()

    workflow = VisualObjectWorkflow(pil_image, detected_objects)
    workflow_by_image[image] = workflow
    add_cropped_objects(workflow, image, prompt)

    if display_results:
        display_detected_objects(workflow)


def get_pil_image_and_part(
    image: InputImage,
    model: MultimodalModel,
    media_resolution: PartMediaResolutionLevel | None,
) -> tuple[PIL_Image, Part]:
    if isinstance(image, Path):
        image_bytes = image.read_bytes()
    else:
        headers = {"User-Agent": "Mozilla/5.0"}
        req = urllib.request.Request(image, headers=headers)
        with urllib.request.urlopen(req, timeout=10) as response:
            image_bytes = response.read()

    pil_image = PIL.Image.open(io.BytesIO(image_bytes))
    content_part = Part.from_bytes(
        data=image_bytes,
        mime_type="image/*",
        media_resolution=media_resolution,
    )

    return pil_image, content_part


def get_object_detection_config(model: Model) -> GenerateContentConfig:
    # Low randomness for more determinism
    return GenerateContentConfig(
        temperature=0.0,
        top_p=0.0,
        seed=42,
        response_mime_type="application/json",
        response_schema=DetectedObjects,
        thinking_config=get_thinking_config(model),
    )


def add_cropped_objects(
    workflow: VisualObjectWorkflow,
    input: InputImage,
    prompt: str,
    crop_margin: int = CROP_MARGIN_PX,
) -> None:
    cropped_images: list[PIL_Image] = []
    obj_count = len(workflow.detected_objects)
    for obj_order, obj in enumerate(workflow.detected_objects, 1):
        cropped_image, _ = extract_object_image(workflow.source_image, obj, crop_margin)
        cropped_images.append(cropped_image)
        save_workflow_image(
            WorkflowStep.SOURCE,
            WorkflowStep.CROPPED,
            input,
            obj_order,
            obj_count,
            cropped_image,
            dict(prompt=prompt, crop_margin=str(crop_margin)),
        )
    workflow.images_by_step[WorkflowStep.CROPPED] = cropped_images


def extract_object_image(
    image: PIL_Image,
    obj: DetectedObject,
    margin: int = 0,
) -> tuple[PIL_Image, tuple[int, int, int, int]]:
    def clamp(coord: int, dim: int) -> int:
        return min(max(coord, 0), dim)

    x1, y1, x2, y2 = obj.box_2d
    w, h = image.size
    if margin != 0:
        x1, x2 = clamp(x1 - margin, w), clamp(x2 + margin, w)
        y1, y2 = clamp(y1 - margin, h), clamp(y2 + margin, h)

    box = (x1, y1, x2, y2)
    object_image = image.crop(box)

    return object_image, box


def save_workflow_image(
    source_step: WorkflowStep,
    target_step: WorkflowStep,
    input_image: InputImage,
    obj_order: int,
    obj_count: int,
    target_image: PIL_Image | None,
    image_info: dict[str, str] | None = None,
) -> None:
    if not SAVE_GENERATED_IMAGES or target_image is None:
        return
    if not OUTPUT_IMAGES_PATH.is_dir():
        OUTPUT_IMAGES_PATH.mkdir(parents=True)
    time_str = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
    try:
        filename = f"{Source(input_image).name}_"
    except ValueError:
        filename = ""
    filename += f"{obj_order}o{obj_count}_{source_step}_{target_step}_{time_str}.png"
    image_path = OUTPUT_IMAGES_PATH.joinpath(filename)
    params = {}
    if image_info:
        png_info = PngInfo()
        for k, v in image_info.items():
            png_info.add_text(k, v)
        params.update(pnginfo=png_info)
    target_image.save(image_path, **params)


# Matplotlib
FIGURE_FG_COLOR = "#F1F3F4"
FIGURE_BG_COLOR = "#202124"
EDGE_COLOR = "#80868B"
rcParams = {
    "figure.dpi": 300,
    "text.color": FIGURE_FG_COLOR,
    "figure.edgecolor": FIGURE_FG_COLOR,
    "axes.titlecolor": FIGURE_FG_COLOR,
    "axes.edgecolor": FIGURE_FG_COLOR,
    "xtick.color": FIGURE_FG_COLOR,
    "ytick.color": FIGURE_FG_COLOR,
    "figure.facecolor": FIGURE_BG_COLOR,
    "axes.edgecolor": EDGE_COLOR,
    "xtick.bottom": False,
    "xtick.top": False,
    "ytick.left": False,
    "ytick.right": False,
    "xtick.labelbottom": False,
    "ytick.labelleft": False,
}
plt.rcParams.update(rcParams)


def display_image_source_info(image: InputImage) -> None:
    def get_image_info_md() -> str:
        if image not in Source:
            return f"[[Source Image]({image})]"
        source = Source(image)
        metadata = metadata_by_source.get(source)
        if not metadata:
            return f"[[Source Image]({source.value})]"
        parts = [
            f"[Source Image]({source.value})",
            f"[Source Page]({metadata.webpage_url})",
            metadata.title,
            metadata.credit_line,
        ]
        separator = "•"
        inner_info = f" {separator} ".join(parts)
        return f"{separator} {inner_info} {separator}"

    def yield_md_rows() -> Iterator[str]:
        horizontal_line = "---"
        image_info = get_image_info_md()
        yield horizontal_line
        yield f"_{image_info}_"
        yield horizontal_line

    display_markdown(f"{chr(10)}{chr(10)}".join(yield_md_rows()))


def display_detected_objects(workflow: VisualObjectWorkflow) -> None:
    source_image = workflow.source_image
    detected_objects = PIL.Image.new("RGB", source_image.size, "white")
    for obj in workflow.detected_objects:
        obj_image, box = extract_object_image(source_image, obj)
        detected_objects.paste(obj_image, (box[0], box[1]))

    _, (ax1, ax2) = plt.subplots(1, 2, layout="compressed")
    ax1.imshow(source_image)
    ax2.imshow(detected_objects)

    disable_colab_cell_scrollbar()
    plt.show()


print("✅ Object detection helpers defined")

from collections import defaultdict


def display_detection_data(source: Source, show_consolidated: bool = False) -> None:
    def string_with_visible_linebreaks(s: str) -> str:
        return f'''"{s.replace(chr(10), "↩️")}"'''

    def yield_md_rows_consolidated(workflow: VisualObjectWorkflow) -> Iterator[str]:
        yield "| label | count | captions |"
        yield "| :--- | ---: | :--- |"
        stats = defaultdict(list)
        for obj in workflow.detected_objects:
            stats[obj.label].append(string_with_visible_linebreaks(obj.caption))
        for label, captions in stats.items():
            count = len(captions)
            label_captions = " • ".join(sorted(captions))
            yield f"| {label} | {count} | {label_captions} |"

    def yield_md_rows_with_bbox(workflow: VisualObjectWorkflow) -> Iterator[str]:
        yield "| box_2d | label | caption |"
        yield "| :--- | :--- | :--- |"
        for obj in workflow.detected_objects:
            yield f"| {obj.box_2d} | {obj.label} | {string_with_visible_linebreaks(obj.caption)} |"

    workflow = workflow_by_image.get(source)
    if workflow is None:
        print(f'❌ No detection for source "{source.name}"')
        return
    md_rows = list(
        yield_md_rows_consolidated(workflow)
        if show_consolidated
        else yield_md_rows_with_bbox(workflow)
    )
    display_image_source_info(source)
    display_markdown(chr(10).join(md_rows))

from typing import Protocol


class ObjectEditingFunction(Protocol):
    def __call__(
        self,
        image: InputImage,
        prompt: str | None = None,
        model: ImageModel | None = None,
        config: GenerateContentConfig | None = None,
        display_results: bool = True,
    ) -> None: ...


SourceTargetSteps = tuple[WorkflowStep, WorkflowStep]
registered_functions: dict[SourceTargetSteps, ObjectEditingFunction] = {}

DEFAULT_EDITING_CONFIG = GenerateContentConfig(response_modalities=["IMAGE"])
EMPTY_IMAGE = PIL.Image.new("1", (1, 1), "white")


def object_editing_function(
    default_prompt: str,
    source_step: WorkflowStep,
    target_step: WorkflowStep,
    default_model: ImageModel = ImageModel.DEFAULT,
    default_config: GenerateContentConfig = DEFAULT_EDITING_CONFIG,
) -> ObjectEditingFunction:
    def editing_function(
        image: InputImage,
        prompt: str | None = default_prompt,
        model: ImageModel | None = default_model,
        config: GenerateContentConfig | None = default_config,
        display_results: bool = True,
    ) -> None:
        workflow, source_images = get_workflow_and_step_images(image, source_step)
        if prompt is None:
            prompt = default_prompt
        prompt = prompt.strip()
        if model is None:
            model = default_model
        # Note: "config is None" is valid and will use the model endpoint default config

        target_images: list[PIL_Image] = []
        display_image_source_info(image)
        obj_count = len(source_images)
        for obj_order, source_image in enumerate(source_images, 1):
            target_image = generate_image([source_image], prompt, model, config)
            save_workflow_image(
                source_step,
                target_step,
                image,
                obj_order,
                obj_count,
                target_image,
                dict(prompt=prompt),
            )
            target_images.append(target_image if target_image else EMPTY_IMAGE)

        workflow.images_by_step[target_step] = target_images
        if display_results:
            display_sources_and_targets(workflow, source_step, target_step)

    registered_functions[(source_step, target_step)] = editing_function

    return editing_function


def get_workflow_and_step_images(
    image: InputImage,
    step: WorkflowStep,
) -> tuple[VisualObjectWorkflow, list[PIL_Image]]:
    # Objects detected?
    if image not in workflow_by_image:
        detect_objects(image, display_results=False)
    workflow = workflow_by_image.get(image)
    assert workflow is not None

    # Workflow step objects? (single level, could be extended to a dynamical graph)
    operation = (WorkflowStep.CROPPED, step)
    if step not in workflow.images_by_step and operation in registered_functions:
        source_function = registered_functions[operation]
        source_function(image, display_results=False)

    # Source images
    source_images = workflow.images_by_step.get(step)
    assert source_images is not None

    return workflow, source_images


def display_sources_and_targets(
    workflow: VisualObjectWorkflow,
    source_step: WorkflowStep,
    target_step: WorkflowStep,
) -> None:
    source_images = workflow.images_by_step[source_step]
    target_images = workflow.images_by_step[target_step]
    if not source_images:
        print("❌ No images to display")
        return

    fig = plt.figure(layout="compressed")
    if horizontal := (len(source_images) >= 2):
        rows, cols = 2, len(source_images)
    else:
        rows, cols = len(source_images), 2
    gs = fig.add_gridspec(rows, cols)

    for i, (source_image, target_image) in enumerate(
        zip(source_images, target_images, strict=True)
    ):
        for dim, image in enumerate([source_image, target_image]):
            grid_spec = gs[dim, i] if horizontal else gs[i, dim]
            ax = fig.add_subplot(grid_spec)
            ax.set_axis_off()
            ax.imshow(image)

    disable_colab_cell_scrollbar()
    plt.show()


print("✅ Object editing helpers defined")