#!/usr/bin/env python3 """ Standalone CV Synthetic Data Engine This module implements an API-first synthetic data pipeline for few-shot object conditioning, prompt-driven synthetic frame generation, Grounding DINO zero-shot annotation, and dataset compilation for YOLO or COCO-style training workflows. The script is intentionally self-contained. It can be used directly from a GPU host, wrapped by Modal serverless functions, or launched inside RunPod. """ from __future__ import annotations import argparse import gc import json import logging import math import os import random import shutil import sys import time from dataclasses import asdict, dataclass, field from pathlib import Path from typing import Any, Dict, Iterable, List, Literal, Optional, Sequence, Tuple, Union import numpy as np from PIL import Image, ImageEnhance, ImageFilter, ImageOps try: import cv2 except Exception as exc: # pragma: no cover - dependency guard cv2 = None _CV2_IMPORT_ERROR = exc else: _CV2_IMPORT_ERROR = None try: import torch import torch.nn.functional as F from torch.utils.data import DataLoader, Dataset except Exception as exc: # pragma: no cover - dependency guard torch = None F = None DataLoader = object Dataset = object _TORCH_IMPORT_ERROR = exc else: _TORCH_IMPORT_ERROR = None try: from diffusers import DDPMScheduler, StableDiffusionPipeline, StableDiffusionXLPipeline from peft.utils import get_peft_model_state_dict except Exception as exc: # pragma: no cover - dependency guard DDPMScheduler = None StableDiffusionPipeline = None StableDiffusionXLPipeline = None get_peft_model_state_dict = None _DIFFUSERS_IMPORT_ERROR = exc else: _DIFFUSERS_IMPORT_ERROR = None try: from peft import LoraConfig except Exception as exc: # pragma: no cover - dependency guard LoraConfig = None _PEFT_IMPORT_ERROR = exc else: _PEFT_IMPORT_ERROR = None try: from transformers import AutoModelForZeroShotObjectDetection, AutoProcessor except Exception as exc: # pragma: no cover - dependency guard AutoModelForZeroShotObjectDetection = None AutoProcessor = None _TRANSFORMERS_IMPORT_ERROR = exc else: _TRANSFORMERS_IMPORT_ERROR = None try: import modal except Exception: # pragma: no cover - optional platform integration modal = None LOGGER = logging.getLogger("synthetic_cv_pipeline") LabelFormat = Literal["yolo", "coco"] ImageLike = Union[str, Path, Image.Image, np.ndarray] @dataclass class TrainingConfig: """Configuration for high-velocity few-shot LoRA optimization.""" pretrained_model: str = "runwayml/stable-diffusion-v1-5" output_dir: str = "conditioned_lora" instance_token: str = "sksobj" resolution: int = 512 train_steps: int = 180 validation_interval: int = 20 patience: int = 4 min_delta: float = 0.0025 learning_rate: float = 1e-4 batch_size: int = 1 gradient_accumulation_steps: int = 1 rank: int = 8 seed: int = 1337 mixed_precision: Literal["fp16", "bf16", "no"] = "fp16" negative_prompt: str = "low quality, blurry, warped object, extra object, text, watermark" num_validation_images: int = 1 train_text_encoder_lora: bool = False max_grad_norm: float = 1.0 @dataclass class SynthesisConfig: """Configuration for conditioned batch synthesis.""" prompt: str target_object: str output_dir: str = "output_batch" lora_dir: Optional[str] = "conditioned_lora" pretrained_model: str = "runwayml/stable-diffusion-v1-5" num_images: int = 24 batch_size: int = 1 width: int = 512 height: int = 512 inference_steps: int = 32 guidance_scale: float = 7.0 lora_scale: float = 0.85 seed: int = 1337 label_format: LabelFormat = "yolo" class_id: int = 0 category_id: int = 1 detector_model: str = "IDEA-Research/grounding-dino-tiny" detection_threshold: float = 0.35 text_threshold: float = 0.25 debug_preview_count: int = 5 negative_prompt: str = "low quality, blurry, duplicate object, malformed, noisy, text, watermark" @dataclass class DetectionRecord: """Normalized internal representation of one detector output.""" image_id: int label: str score: float box_xyxy: Tuple[float, float, float, float] width: int height: int def clipped(self) -> "DetectionRecord": xmin, ymin, xmax, ymax = self.box_xyxy xmin = float(max(0.0, min(xmin, self.width - 1))) ymin = float(max(0.0, min(ymin, self.height - 1))) xmax = float(max(0.0, min(xmax, self.width - 1))) ymax = float(max(0.0, min(ymax, self.height - 1))) if xmax < xmin: xmin, xmax = xmax, xmin if ymax < ymin: ymin, ymax = ymax, ymin return DetectionRecord( image_id=self.image_id, label=self.label, score=self.score, box_xyxy=(xmin, ymin, xmax, ymax), width=self.width, height=self.height, ) def configure_logging(level: str = "INFO") -> None: numeric_level = getattr(logging, level.upper(), logging.INFO) logging.basicConfig( level=numeric_level, format="%(asctime)s | %(levelname)s | %(name)s | %(message)s", datefmt="%Y-%m-%d %H:%M:%S", ) def require_dependency(name: str, import_error: Optional[BaseException]) -> None: if import_error is not None: raise RuntimeError( f"Required dependency '{name}' could not be imported. Install the expected GPU stack " f"before running this pipeline. Original error: {import_error}" ) from import_error def resolve_device() -> str: require_dependency("torch", _TORCH_IMPORT_ERROR) if torch.cuda.is_available(): return "cuda" if hasattr(torch.backends, "mps") and torch.backends.mps.is_available(): return "mps" return "cpu" def clear_vram() -> None: gc.collect() if torch is not None and torch.cuda.is_available(): torch.cuda.empty_cache() torch.cuda.ipc_collect() def load_source_images(images: Sequence[ImageLike]) -> List[Image.Image]: loaded: List[Image.Image] = [] for idx, item in enumerate(images): try: if isinstance(item, Image.Image): image = item.convert("RGB") elif isinstance(item, np.ndarray): arr = item if arr.ndim == 2: arr = np.stack([arr] * 3, axis=-1) if arr.shape[-1] == 4: arr = arr[..., :3] image = Image.fromarray(arr.astype(np.uint8)).convert("RGB") else: path = Path(item).expanduser().resolve() if not path.exists(): raise FileNotFoundError(f"Source image does not exist: {path}") image = Image.open(path).convert("RGB") loaded.append(image) except Exception as exc: raise ValueError(f"Failed to load source image at index {idx}: {exc}") from exc if not loaded: raise ValueError("At least one source target object image is required.") return loaded def load_image_paths(directory: Union[str, Path]) -> List[Path]: root = Path(directory).expanduser().resolve() if not root.exists() or not root.is_dir(): raise FileNotFoundError(f"Source image directory not found: {root}") allowed = {".jpg", ".jpeg", ".png", ".webp", ".bmp"} paths = sorted(path for path in root.iterdir() if path.suffix.lower() in allowed) if not paths: raise FileNotFoundError(f"No supported images found in {root}") return paths def pil_to_tensor(image: Image.Image, resolution: int) -> "torch.Tensor": image = ImageOps.exif_transpose(image).convert("RGB") image = ImageOps.fit(image, (resolution, resolution), method=Image.Resampling.LANCZOS) arr = np.asarray(image).astype(np.float32) / 127.5 - 1.0 tensor = torch.from_numpy(arr).permute(2, 0, 1) return tensor class FewShotImageDataset(Dataset): """A deterministic plus stochastic few-shot image dataset for LoRA training.""" def __init__(self, images: Sequence[Image.Image], prompt: str, resolution: int, length: int = 2048) -> None: self.images = list(images) self.prompt = prompt self.resolution = resolution self.length = max(length, len(self.images)) def __len__(self) -> int: return self.length def _augment(self, image: Image.Image) -> Image.Image: image = ImageOps.exif_transpose(image).convert("RGB") if random.random() < 0.5: image = ImageOps.mirror(image) brightness = random.uniform(0.82, 1.18) contrast = random.uniform(0.85, 1.15) saturation = random.uniform(0.82, 1.2) image = ImageEnhance.Brightness(image).enhance(brightness) image = ImageEnhance.Contrast(image).enhance(contrast) image = ImageEnhance.Color(image).enhance(saturation) angle = random.uniform(-8, 8) image = image.rotate(angle, resample=Image.Resampling.BICUBIC, expand=False, fillcolor=(127, 127, 127)) return image def __getitem__(self, index: int) -> Dict[str, Any]: image = self.images[index % len(self.images)] return {"pixel_values": pil_to_tensor(self._augment(image), self.resolution), "prompt": self.prompt} def compute_structural_loss(candidate: Image.Image, references: Sequence[Image.Image], resolution: int) -> float: """ Compute a lightweight structural validation loss from edge maps and luminance. The validation signal intentionally emphasizes shape retention instead of exact background fidelity. This lets early stopping preserve target object features while avoiding overfitting to source lighting and environment. """ require_dependency("opencv-python", _CV2_IMPORT_ERROR) candidate_gray = np.asarray(ImageOps.fit(candidate.convert("L"), (resolution, resolution), Image.Resampling.LANCZOS)) candidate_edges = cv2.Canny(candidate_gray, 80, 160).astype(np.float32) / 255.0 candidate_luma = candidate_gray.astype(np.float32) / 255.0 best_loss = float("inf") for ref in references: ref_gray = np.asarray(ImageOps.fit(ref.convert("L"), (resolution, resolution), Image.Resampling.LANCZOS)) ref_edges = cv2.Canny(ref_gray, 80, 160).astype(np.float32) / 255.0 ref_luma = ref_gray.astype(np.float32) / 255.0 edge_loss = float(np.mean((candidate_edges - ref_edges) ** 2)) luma_loss = float(np.mean((candidate_luma - ref_luma) ** 2)) best_loss = min(best_loss, 0.75 * edge_loss + 0.25 * luma_loss) return best_loss def create_prompt_variation(base_prompt: str, target_object: str, instance_token: str, index: int) -> str: camera_angles = [ "front three-quarter view", "low angle macro shot", "high angle inspection view", "side profile perspective", "telephoto compressed perspective", "wide-angle close pass", ] illumination = [ "soft diffuse overcast lighting", "hard rim light with long shadows", "cool fluorescent industrial illumination", "warm golden hour side light", "dramatic backlight and controlled reflections", "mixed practical lights with subtle glare", ] materials = [ "mild specular reflections", "matte surface response", "gloss highlights on nearby surfaces", "wet floor reflections", "dusty atmospheric scattering", "clean studio-grade clarity", ] distance = [ "object occupying 12 percent of the frame", "object occupying 25 percent of the frame", "object occupying 40 percent of the frame", "object in the near foreground", "object at medium distance", "object partially framed by environmental structures", ] occlusion = [ "unoccluded target object", "subtle foreground occlusion at one edge", "partial shadow crossing the target object", "thin cable-like occluder in foreground", "minor motion blur in the environment only", "clean silhouette with no occlusion", ] rng = random.Random(index * 7919 + len(base_prompt)) descriptors = [ rng.choice(camera_angles), rng.choice(illumination), rng.choice(materials), rng.choice(distance), rng.choice(occlusion), ] descriptor_text = ", ".join(descriptors) return f"a photo of {instance_token} {target_object} in {base_prompt}, {descriptor_text}, realistic, high detail" def encode_prompt_for_training(pipe: Any, prompt: Union[str, List[str]], device: str) -> "torch.Tensor": tokens = pipe.tokenizer( prompt, padding="max_length", max_length=pipe.tokenizer.model_max_length, truncation=True, return_tensors="pt", ) input_ids = tokens.input_ids.to(device) return pipe.text_encoder(input_ids)[0] class FewShotLoRATrainer: """Few-shot LoRA trainer with early stopping on validation structural loss.""" def __init__(self, config: TrainingConfig) -> None: self.config = config self.device = resolve_device() self.dtype = self._resolve_dtype() self.pipe: Optional[Any] = None def _resolve_dtype(self) -> "torch.dtype": if torch is None: raise RuntimeError("PyTorch is required for training.") if self.config.mixed_precision == "bf16" and torch.cuda.is_available() and torch.cuda.is_bf16_supported(): return torch.bfloat16 if self.config.mixed_precision == "fp16" and self.device == "cuda": return torch.float32 return torch.float32 def _load_pipeline(self) -> Any: require_dependency("diffusers", _DIFFUSERS_IMPORT_ERROR) require_dependency("peft", _PEFT_IMPORT_ERROR) LOGGER.info("Loading diffusion pipeline for LoRA training: %s", self.config.pretrained_model) try: pipe = StableDiffusionPipeline.from_pretrained( self.config.pretrained_model, torch_dtype=self.dtype, safety_checker=None, requires_safety_checker=False, ) pipe.scheduler = DDPMScheduler.from_config(pipe.scheduler.config) pipe.to(self.device) pipe.vae.requires_grad_(False) pipe.text_encoder.requires_grad_(False) pipe.unet.requires_grad_(False) lora_config = LoraConfig( r=self.config.rank, lora_alpha=self.config.rank, init_lora_weights="gaussian", target_modules=["to_k", "to_q", "to_v", "to_out.0"], ) pipe.unet.add_adapter(lora_config) if self.config.train_text_encoder_lora: text_lora_config = LoraConfig( r=self.config.rank, lora_alpha=self.config.rank, init_lora_weights="gaussian", target_modules=["q_proj", "k_proj", "v_proj", "out_proj"], ) pipe.text_encoder.add_adapter(text_lora_config) if hasattr(pipe, "enable_xformers_memory_efficient_attention"): try: pipe.enable_xformers_memory_efficient_attention() LOGGER.info("Enabled xFormers memory-efficient attention.") except Exception as exc: LOGGER.warning("Could not enable xFormers attention: %s", exc) if hasattr(pipe, "enable_attention_slicing"): pipe.enable_attention_slicing() self.pipe = pipe return pipe except torch.cuda.OutOfMemoryError as exc: clear_vram() raise RuntimeError("CUDA VRAM exhausted while loading the diffusion training pipeline.") from exc except Exception as exc: clear_vram() raise RuntimeError(f"Failed to load diffusion training pipeline: {exc}") from exc def train(self, source_images: Sequence[ImageLike], target_object: str) -> Path: images = load_source_images(source_images) random.seed(self.config.seed) np.random.seed(self.config.seed) torch.manual_seed(self.config.seed) if torch.cuda.is_available(): torch.cuda.manual_seed_all(self.config.seed) output_dir = Path(self.config.output_dir).expanduser().resolve() output_dir.mkdir(parents=True, exist_ok=True) prompt = f"a photo of {self.config.instance_token} {target_object}" dataset = FewShotImageDataset(images, prompt=prompt, resolution=self.config.resolution) loader = DataLoader(dataset, batch_size=self.config.batch_size, shuffle=True, num_workers=0) iterator = iter(loader) pipe = self._load_pipeline() trainable_params = [p for p in pipe.unet.parameters() if p.requires_grad] if self.config.train_text_encoder_lora: trainable_params += [p for p in pipe.text_encoder.parameters() if p.requires_grad] if not trainable_params: raise RuntimeError("No trainable LoRA parameters were registered. Check PEFT and Diffusers versions.") optimizer = torch.optim.AdamW(trainable_params, lr=self.config.learning_rate, betas=(0.9, 0.999), weight_decay=0.01) scaler_enabled = self.dtype == torch.float32 and self.device == "cuda" scaler = torch.cuda.amp.GradScaler(enabled=False) best_structural_loss = float("inf") patience_counter = 0 global_step = 0 running_loss = 0.0 start = time.time() pipe.unet.train() if self.config.train_text_encoder_lora: pipe.text_encoder.train() LOGGER.info("Starting LoRA optimization with %d source images for target '%s'.", len(images), target_object) while global_step < self.config.train_steps: try: batch = next(iterator) except StopIteration: iterator = iter(loader) batch = next(iterator) pixel_values = batch["pixel_values"].to(device=self.device, dtype=self.dtype) prompts = list(batch["prompt"]) with torch.no_grad(): latents = pipe.vae.encode(pixel_values).latent_dist.sample() latents = latents * pipe.vae.config.scaling_factor noise = torch.randn_like(latents) timesteps = torch.randint( 0, pipe.scheduler.config.num_train_timesteps, (latents.shape[0],), device=self.device, dtype=torch.long, ) noisy_latents = pipe.scheduler.add_noise(latents, noise, timesteps) encoder_hidden_states = encode_prompt_for_training(pipe, prompts, self.device).to(dtype=self.dtype) try: with torch.amp.autocast('cuda', dtype=torch.float32): model_pred = pipe.unet(noisy_latents, timesteps, encoder_hidden_states).sample target = noise if getattr(pipe.scheduler.config, "prediction_type", None) == "v_prediction": target = pipe.scheduler.get_velocity(latents, noise, timesteps) loss = F.mse_loss(model_pred.float(), target.float(), reduction="mean") loss = loss / self.config.gradient_accumulation_steps scaler.scale(loss).backward() except torch.cuda.OutOfMemoryError as exc: clear_vram() raise RuntimeError("CUDA VRAM exhausted during LoRA optimization. Reduce resolution, rank, or batch size.") from exc running_loss += float(loss.detach().cpu().item()) if (global_step + 1) % self.config.gradient_accumulation_steps == 0: scaler.unscale_(optimizer) torch.nn.utils.clip_grad_norm_(trainable_params, self.config.max_grad_norm) scaler.step(optimizer) scaler.update() optimizer.zero_grad(set_to_none=True) global_step += 1 if global_step % max(1, self.config.validation_interval) == 0 or global_step == self.config.train_steps: structural_loss = self._validate_structural_loss(pipe, images, target_object) avg_train_loss = running_loss / max(1, self.config.validation_interval) running_loss = 0.0 LOGGER.info( "step=%d/%d train_loss=%.6f validation_structural_loss=%.6f best=%.6f elapsed=%.1fs", global_step, self.config.train_steps, avg_train_loss, structural_loss, best_structural_loss, time.time() - start, ) if structural_loss + self.config.min_delta < best_structural_loss: best_structural_loss = structural_loss patience_counter = 0 self._save_lora(pipe, output_dir) LOGGER.info("Saved improved LoRA checkpoint to %s", output_dir) else: patience_counter += 1 if patience_counter >= self.config.patience: LOGGER.info( "Early stopping triggered at step %d after %d stagnant validations.", global_step, patience_counter, ) break self._save_lora(pipe, output_dir) metadata = { "target_object": target_object, "instance_token": self.config.instance_token, "pretrained_model": self.config.pretrained_model, "best_structural_loss": best_structural_loss, "steps_completed": global_step, "training_config": asdict(self.config), } (output_dir / "conditioning_metadata.json").write_text(json.dumps(metadata, indent=2), encoding="utf-8") LOGGER.info("Training complete. LoRA artifacts are available in %s", output_dir) return output_dir def _validate_structural_loss(self, pipe: Any, references: Sequence[Image.Image], target_object: str) -> float: pipe.unet.eval() if self.config.train_text_encoder_lora: pipe.text_encoder.eval() prompt = f"a centered studio product photo of {self.config.instance_token} {target_object}, neutral background, crisp outline" generator = torch.Generator(device=self.device).manual_seed(self.config.seed + 17) losses: List[float] = [] try: with torch.no_grad(): generated = pipe( prompt=prompt, negative_prompt=self.config.negative_prompt, num_images_per_prompt=self.config.num_validation_images, num_inference_steps=18, guidance_scale=6.0, height=self.config.resolution, width=self.config.resolution, generator=generator, ).images for image in generated: losses.append(compute_structural_loss(image, references, self.config.resolution)) except torch.cuda.OutOfMemoryError as exc: clear_vram() raise RuntimeError("CUDA VRAM exhausted during structural validation.") from exc finally: pipe.unet.train() if self.config.train_text_encoder_lora: pipe.text_encoder.train() if not losses: return float("inf") return float(np.mean(losses)) def _save_lora(self, pipe: Any, output_dir: Path) -> None: output_dir.mkdir(parents=True, exist_ok=True) if hasattr(pipe, "save_lora_weights") and get_peft_model_state_dict is not None: save_kwargs: Dict[str, Any] = {"save_directory": str(output_dir), "unet_lora_layers": get_peft_model_state_dict(pipe.unet)} if self.config.train_text_encoder_lora: save_kwargs["text_encoder_lora_layers"] = get_peft_model_state_dict(pipe.text_encoder) pipe.save_lora_weights(**save_kwargs) elif hasattr(pipe.unet, "save_pretrained"): pipe.unet.save_pretrained(str(output_dir / "unet_lora")) else: raise RuntimeError("The loaded pipeline cannot save LoRA weights with the installed Diffusers version.") class ParametricSynthesizer: """Prompt-conditioned synthetic frame generator with systematic visual variation.""" def __init__(self, config: SynthesisConfig, instance_token: str = "sksobj") -> None: self.config = config self.instance_token = instance_token self.device = resolve_device() self.dtype = torch.float32 if self.device == "cuda" else torch.float32 self.pipe: Optional[Any] = None def _load_pipeline(self) -> Any: require_dependency("diffusers", _DIFFUSERS_IMPORT_ERROR) LOGGER.info("Loading synthesis pipeline: %s", self.config.pretrained_model) try: is_sdxl = "xl" in self.config.pretrained_model.lower() or "sdxl" in self.config.pretrained_model.lower() pipeline_cls = StableDiffusionXLPipeline if is_sdxl else StableDiffusionPipeline load_kwargs: Dict[str, Any] = {"torch_dtype": self.dtype} if not is_sdxl: load_kwargs.update({"safety_checker": None, "requires_safety_checker": False}) pipe = pipeline_cls.from_pretrained(self.config.pretrained_model, **load_kwargs) pipe.to(self.device) if hasattr(pipe, "enable_attention_slicing"): pipe.enable_attention_slicing() if hasattr(pipe, "enable_vae_slicing"): pipe.enable_vae_slicing() if self.config.lora_dir: lora_path = Path(self.config.lora_dir).expanduser().resolve() if lora_path.exists(): pipe.load_lora_weights(str(lora_path)) if hasattr(pipe, "set_adapters"): try: pipe.set_adapters(["default_0"], adapter_weights=[self.config.lora_scale]) except Exception: LOGGER.debug("Adapter weighting API unavailable or adapter name differs; using loaded LoRA default scale.") LOGGER.info("Loaded LoRA weights from %s", lora_path) else: raise FileNotFoundError(f"Configured LoRA directory does not exist: {lora_path}") self.pipe = pipe return pipe except torch.cuda.OutOfMemoryError as exc: clear_vram() raise RuntimeError("CUDA VRAM exhausted while loading the synthesis pipeline.") from exc except Exception as exc: clear_vram() raise RuntimeError(f"Failed to load synthesis pipeline: {exc}") from exc def generate(self) -> List[Tuple[Path, Image.Image, str]]: pipe = self.pipe or self._load_pipeline() output_root = Path(self.config.output_dir).expanduser().resolve() image_dir = output_root / "images" label_dir = output_root / "labels" image_dir.mkdir(parents=True, exist_ok=True) label_dir.mkdir(parents=True, exist_ok=True) generated_records: List[Tuple[Path, Image.Image, str]] = [] LOGGER.info("Generating %d synthetic frames into %s", self.config.num_images, image_dir) for start_idx in range(0, self.config.num_images, self.config.batch_size): current_batch = min(self.config.batch_size, self.config.num_images - start_idx) prompts = [ create_prompt_variation(self.config.prompt, self.config.target_object, self.instance_token, start_idx + i) for i in range(current_batch) ] generators = [torch.Generator(device=self.device).manual_seed(self.config.seed + start_idx + i) for i in range(current_batch)] try: with torch.no_grad(): result = pipe( prompt=prompts, negative_prompt=[self.config.negative_prompt] * current_batch, width=self.config.width, height=self.config.height, num_inference_steps=self.config.inference_steps, guidance_scale=self.config.guidance_scale, generator=generators, ) except torch.cuda.OutOfMemoryError as exc: clear_vram() raise RuntimeError("CUDA VRAM exhausted during synthesis. Lower batch size, resolution, or steps.") from exc for local_idx, image in enumerate(result.images): image_id = start_idx + local_idx post_image = self._postprocess_variation(image, image_id) image_path = image_dir / f"synthetic_{image_id:06d}.jpg" post_image.save(image_path, quality=95) generated_records.append((image_path, post_image, prompts[local_idx])) LOGGER.debug("Generated %s with prompt: %s", image_path.name, prompts[local_idx]) LOGGER.info("Generated %d images.", len(generated_records)) return generated_records def _postprocess_variation(self, image: Image.Image, index: int) -> Image.Image: rng = random.Random(self.config.seed + index * 1543) image = image.convert("RGB") if rng.random() < 0.35: overlay = Image.new("RGB", image.size, (255, 255, 255)) alpha = rng.uniform(0.015, 0.06) image = Image.blend(image, overlay, alpha) if rng.random() < 0.35: image = ImageEnhance.Brightness(image).enhance(rng.uniform(0.88, 1.12)) if rng.random() < 0.35: image = ImageEnhance.Contrast(image).enhance(rng.uniform(0.9, 1.15)) if rng.random() < 0.25: image = image.filter(ImageFilter.GaussianBlur(radius=rng.uniform(0.0, 0.45))) return image class GroundingDINOLabeler: """Grounding DINO wrapper for zero-shot bounding-box extraction.""" def __init__(self, model_id: str, box_threshold: float, text_threshold: float) -> None: require_dependency("transformers", _TRANSFORMERS_IMPORT_ERROR) require_dependency("torch", _TORCH_IMPORT_ERROR) self.model_id = model_id self.box_threshold = box_threshold self.text_threshold = text_threshold self.device = resolve_device() LOGGER.info("Loading Grounding DINO detector: %s", model_id) try: self.processor = AutoProcessor.from_pretrained(model_id) self.model = AutoModelForZeroShotObjectDetection.from_pretrained(model_id).to(self.device) self.model.eval() except torch.cuda.OutOfMemoryError as exc: clear_vram() raise RuntimeError("CUDA VRAM exhausted while loading Grounding DINO.") from exc except Exception as exc: clear_vram() raise RuntimeError(f"Failed to load Grounding DINO model '{model_id}': {exc}") from exc def detect(self, image: Image.Image, query: str, image_id: int) -> List[DetectionRecord]: text_labels = [[query]] width, height = image.size try: inputs = self.processor(images=image, text=text_labels, return_tensors="pt").to(self.device) with torch.no_grad(): outputs = self.model(**inputs) results = self.processor.post_process_grounded_object_detection( outputs, inputs.input_ids, threshold=self.box_threshold, text_threshold=self.text_threshold, target_sizes=[(height, width)], )[0] except torch.cuda.OutOfMemoryError as exc: clear_vram() raise RuntimeError("CUDA VRAM exhausted during Grounding DINO inference.") from exc except Exception as exc: raise RuntimeError(f"Grounding DINO inference failed for image_id={image_id}: {exc}") from exc records: List[DetectionRecord] = [] boxes = results.get("boxes", []) scores = results.get("scores", []) labels = results.get("labels", []) for box, score, label in zip(boxes, scores, labels): box_tuple = tuple(float(x) for x in box.detach().cpu().tolist()) score_float = float(score.detach().cpu().item()) if hasattr(score, "detach") else float(score) label_text = str(label) rec = DetectionRecord(image_id=image_id, label=label_text, score=score_float, box_xyxy=box_tuple, width=width, height=height).clipped() xmin, ymin, xmax, ymax = rec.box_xyxy if xmax - xmin >= 2 and ymax - ymin >= 2: records.append(rec) LOGGER.debug("Detector returned %d boxes for image_id=%d", len(records), image_id) return records def convert_detection(record: DetectionRecord, fmt: LabelFormat, class_id: int = 0, category_id: int = 1) -> Union[List[float], Dict[str, Any]]: rec = record.clipped() xmin, ymin, xmax, ymax = rec.box_xyxy box_w = max(0.0, xmax - xmin) box_h = max(0.0, ymax - ymin) if fmt == "yolo": x_center = (xmin + box_w / 2.0) / rec.width y_center = (ymin + box_h / 2.0) / rec.height return [ int(class_id), round(float(x_center), 6), round(float(y_center), 6), round(float(box_w / rec.width), 6), round(float(box_h / rec.height), 6), ] if fmt == "coco": return { "image_id": int(rec.image_id), "category_id": int(category_id), "bbox": [round(float(xmin), 2), round(float(ymin), 2), round(float(box_w), 2), round(float(box_h), 2)], "score": round(float(rec.score), 6), "label": rec.label, } raise ValueError(f"Unsupported label format: {fmt}") def write_label_file(label_dir: Path, image_path: Path, detections: Sequence[DetectionRecord], config: SynthesisConfig) -> Path: label_dir.mkdir(parents=True, exist_ok=True) if config.label_format == "yolo": label_path = label_dir / f"{image_path.stem}.txt" lines = [] for rec in detections: converted = convert_detection(rec, "yolo", class_id=config.class_id) lines.append(" ".join(str(x) for x in converted)) label_path.write_text("\n".join(lines) + ("\n" if lines else ""), encoding="utf-8") return label_path label_path = label_dir / f"{image_path.stem}.json" records = [convert_detection(rec, "coco", category_id=config.category_id) for rec in detections] label_path.write_text(json.dumps(records, indent=2), encoding="utf-8") return label_path def draw_debug_previews( output_dir: Union[str, Path], image_records: Sequence[Tuple[Path, List[DetectionRecord]]], label_format: LabelFormat, count: int = 5, seed: int = 1337, ) -> List[Path]: require_dependency("opencv-python", _CV2_IMPORT_ERROR) output_root = Path(output_dir).expanduser().resolve() if not image_records: LOGGER.warning("No image records available for debug preview generation.") return [] rng = random.Random(seed) sample = list(image_records) rng.shuffle(sample) selected = sample[: min(count, len(sample))] preview_paths: List[Path] = [] for idx, (image_path, detections) in enumerate(selected): img = cv2.imread(str(image_path)) if img is None: LOGGER.warning("OpenCV could not read image for debug preview: %s", image_path) continue for rec in detections: xmin, ymin, xmax, ymax = [int(round(v)) for v in rec.clipped().box_xyxy] cv2.rectangle(img, (xmin, ymin), (xmax, ymax), (40, 220, 40), 2) label = f"{rec.label} {rec.score:.2f}" cv2.putText(img, label, (xmin, max(15, ymin - 6)), cv2.FONT_HERSHEY_SIMPLEX, 0.48, (40, 220, 40), 1, cv2.LINE_AA) preview_path = output_root / f"debug_preview_{idx:02d}.jpg" cv2.imwrite(str(preview_path), img) preview_paths.append(preview_path) LOGGER.info("Saved %s debug preview: %s", label_format.upper(), preview_path) return preview_paths def compile_and_label_outputs(generated: Sequence[Tuple[Path, Image.Image, str]], config: SynthesisConfig) -> Dict[str, Any]: output_root = Path(config.output_dir).expanduser().resolve() image_dir = output_root / "images" label_dir = output_root / "labels" image_dir.mkdir(parents=True, exist_ok=True) label_dir.mkdir(parents=True, exist_ok=True) labeler = GroundingDINOLabeler(config.detector_model, config.detection_threshold, config.text_threshold) image_records: List[Tuple[Path, List[DetectionRecord]]] = [] total_boxes = 0 for image_id, (image_path, image, prompt) in enumerate(generated): detections = labeler.detect(image, query=config.target_object, image_id=image_id) write_label_file(label_dir, image_path, detections, config) image_records.append((image_path, detections)) total_boxes += len(detections) LOGGER.info("Labeled image_id=%d file=%s boxes=%d", image_id, image_path.name, len(detections)) previews = draw_debug_previews(output_root, image_records, config.label_format, config.debug_preview_count, config.seed) manifest = { "output_dir": str(output_root), "image_dir": str(image_dir), "label_dir": str(label_dir), "label_format": config.label_format, "target_object": config.target_object, "num_images": len(generated), "total_boxes": total_boxes, "debug_previews": [str(path) for path in previews], "config": asdict(config), } (output_root / "manifest.json").write_text(json.dumps(manifest, indent=2), encoding="utf-8") LOGGER.info("Compiled labeled dataset in %s with %d boxes.", output_root, total_boxes) return manifest def run_pipeline(source_images: Sequence[ImageLike], train_cfg: TrainingConfig, synth_cfg: SynthesisConfig) -> Dict[str, Any]: LOGGER.info("Starting end-to-end synthetic data pipeline.") trainer = FewShotLoRATrainer(train_cfg) lora_dir = trainer.train(source_images, synth_cfg.target_object) synth_cfg.lora_dir = str(lora_dir) synthesizer = ParametricSynthesizer(synth_cfg, instance_token=train_cfg.instance_token) generated = synthesizer.generate() manifest = compile_and_label_outputs(generated, synth_cfg) LOGGER.info("Pipeline complete: %s", manifest["output_dir"]) return manifest if modal is not None: # pragma: no cover - only active in Modal runtime modal_image = ( modal.Image.debian_slim(python_version="3.11") .pip_install( "torch", "diffusers", "transformers", "accelerate", "opencv-python-headless", "pillow", "peft", "safetensors", ) ) app = modal.App("standalone-cv-synthetic-data-engine", image=modal_image) @app.function(gpu="A10G", timeout=60 * 60 * 3) def modal_run_pipeline(source_dir: str, training_config: Dict[str, Any], synthesis_config: Dict[str, Any]) -> Dict[str, Any]: configure_logging("INFO") paths = load_image_paths(source_dir) train_cfg = TrainingConfig(**training_config) synth_cfg = SynthesisConfig(**synthesis_config) return run_pipeline(paths, train_cfg, synth_cfg) else: app = None def parse_args(argv: Optional[Sequence[str]] = None) -> argparse.Namespace: parser = argparse.ArgumentParser(description="Few-shot synthetic CV data generation pipeline.") parser.add_argument("--source-dir", required=True, help="Directory containing target object source images.") parser.add_argument("--target-object", required=True, help="Exact semantic text string for the target object detector query.") parser.add_argument("--prompt", required=True, help="Background/environment prompt, e.g. 'industrial conveyor belt with reflections'.") parser.add_argument("--output-dir", default="output_batch", help="Output dataset directory.") parser.add_argument("--pretrained-model", default="runwayml/stable-diffusion-v1-5", help="Open diffusion model identifier or local path.") parser.add_argument("--detector-model", default="IDEA-Research/grounding-dino-tiny", help="Grounding DINO model identifier.") parser.add_argument("--label-format", choices=["yolo", "coco"], default="yolo", help="Annotation output format.") parser.add_argument("--num-images", type=int, default=24, help="Number of synthetic images to generate.") parser.add_argument("--train-steps", type=int, default=180, help="Maximum LoRA training steps before early stopping.") parser.add_argument("--resolution", type=int, default=512, help="Training image resolution.") parser.add_argument("--width", type=int, default=512, help="Generated image width.") parser.add_argument("--height", type=int, default=512, help="Generated image height.") parser.add_argument("--batch-size", type=int, default=1, help="Synthesis batch size.") parser.add_argument("--train-batch-size", type=int, default=1, help="LoRA training batch size.") parser.add_argument("--rank", type=int, default=8, help="LoRA rank.") parser.add_argument("--learning-rate", type=float, default=1e-4, help="LoRA learning rate.") parser.add_argument("--seed", type=int, default=1337, help="Random seed.") parser.add_argument("--log-level", default="INFO", help="Logging verbosity.") return parser.parse_args(argv) def main(argv: Optional[Sequence[str]] = None) -> int: args = parse_args(argv) configure_logging(args.log_level) try: source_paths = load_image_paths(args.source_dir) lora_dir = str(Path(args.output_dir).expanduser().resolve() / "conditioned_lora") train_cfg = TrainingConfig( pretrained_model=args.pretrained_model, output_dir=lora_dir, resolution=args.resolution, train_steps=args.train_steps, batch_size=args.train_batch_size, rank=args.rank, learning_rate=args.learning_rate, seed=args.seed, ) synth_cfg = SynthesisConfig( prompt=args.prompt, target_object=args.target_object, output_dir=args.output_dir, lora_dir=lora_dir, pretrained_model=args.pretrained_model, num_images=args.num_images, batch_size=args.batch_size, width=args.width, height=args.height, seed=args.seed, label_format=args.label_format, detector_model=args.detector_model, ) manifest = run_pipeline(source_paths, train_cfg, synth_cfg) LOGGER.info("Final manifest: %s", json.dumps(manifest, indent=2)) return 0 except KeyboardInterrupt: LOGGER.warning("Pipeline interrupted by user.") return 130 except Exception as exc: LOGGER.exception("Pipeline failed: %s", exc) return 1 finally: clear_vram() if __name__ == "__main__": raise SystemExit(main())