| """annotate_crops node — nano-banana (Gemini image generation) for semantic annotation."""
|
| from __future__ import annotations
|
|
|
| import io
|
| from concurrent.futures import ThreadPoolExecutor, as_completed
|
|
|
| from google import genai
|
| from google.genai import types
|
| from PIL import Image
|
|
|
| from config import ANNOTATOR_MODEL, GOOGLE_API_KEY
|
| from prompts.annotator import ANNOTATION_WRAPPER
|
| from state import DrawingReaderState, ImageRef
|
| from tools.image_store import ImageStore
|
|
|
|
|
| def _extract_generated_image(response) -> Image.Image | None:
|
| """Extract the generated image from a Gemini image-generation response."""
|
| for part in response.candidates[0].content.parts:
|
| if part.inline_data is not None:
|
| return Image.open(io.BytesIO(part.inline_data.data))
|
| return None
|
|
|
|
|
| def _annotate_single_crop_sync(
|
| client: genai.Client,
|
| crop_ref: ImageRef,
|
| annotation_prompt: str,
|
| image_store: ImageStore,
|
| ) -> ImageRef | None:
|
| """Annotate one crop using nano-banana (synchronous)."""
|
| crop_bytes = image_store.load_bytes(crop_ref)
|
|
|
| full_prompt = ANNOTATION_WRAPPER.format(annotation_prompt=annotation_prompt)
|
|
|
| response = client.models.generate_content(
|
| model=ANNOTATOR_MODEL,
|
| contents=[
|
| types.Part.from_bytes(data=crop_bytes, mime_type="image/png"),
|
| full_prompt,
|
| ],
|
| config=types.GenerateContentConfig(
|
| response_modalities=["TEXT", "IMAGE"],
|
| ),
|
| )
|
|
|
| annotated_image = _extract_generated_image(response)
|
| if annotated_image is None:
|
| return None
|
|
|
| ref = image_store.save_annotated(crop_ref, annotated_image)
|
| return ref
|
|
|
|
|
| def annotate_crops(state: DrawingReaderState, image_store: ImageStore) -> dict:
|
| """Run nano-banana annotation on crops that need it."""
|
| crop_tasks = state.get("crop_tasks", [])
|
| image_refs = state.get("image_refs", [])
|
|
|
|
|
|
|
|
|
|
|
| crops_needing_annotation: list[tuple[ImageRef, str]] = []
|
|
|
| all_crops = [r for r in image_refs if r["crop_type"] == "crop"]
|
|
|
| recent_crops = all_crops[-len(crop_tasks):] if crop_tasks else []
|
|
|
| for i, task in enumerate(crop_tasks):
|
| if task["annotate"] and task["annotation_prompt"] and i < len(recent_crops):
|
| crops_needing_annotation.append(
|
| (recent_crops[i], task["annotation_prompt"])
|
| )
|
|
|
| if not crops_needing_annotation:
|
| return {"status_message": ["No annotation needed for these crops."]}
|
|
|
| client = genai.Client(api_key=GOOGLE_API_KEY)
|
|
|
|
|
|
|
| results: list[ImageRef | None | Exception] = [None] * len(crops_needing_annotation)
|
|
|
| with ThreadPoolExecutor(max_workers=min(len(crops_needing_annotation), 4)) as pool:
|
| future_to_idx = {}
|
| for i, (ref, prompt) in enumerate(crops_needing_annotation):
|
| future = pool.submit(
|
| _annotate_single_crop_sync, client, ref, prompt, image_store,
|
| )
|
| future_to_idx[future] = i
|
|
|
| for future in as_completed(future_to_idx):
|
| idx = future_to_idx[future]
|
| try:
|
| results[idx] = future.result()
|
| except Exception as e:
|
| results[idx] = e
|
|
|
| annotated_refs: list[ImageRef] = []
|
| errors: list[str] = []
|
| for i, result in enumerate(results):
|
| if isinstance(result, Exception):
|
| errors.append(f"Annotation {i} failed: {result}")
|
| elif result is not None:
|
| annotated_refs.append(result)
|
| else:
|
| errors.append(f"Annotation {i} returned no image")
|
|
|
| status = f"Annotated {len(annotated_refs)} of {len(crops_needing_annotation)} crops."
|
| if errors:
|
| status += f" Issues: {'; '.join(errors)}"
|
|
|
| return {
|
| "image_refs": annotated_refs,
|
| "status_message": [status],
|
| }
|
|
|
