Spaces:
Sleeping
Sleeping
| """Main pipeline: orchestrates segmentation, saliency, and modifications.""" | |
| from __future__ import annotations | |
| import logging | |
| import random | |
| import time | |
| from typing import Callable | |
| import cv2 | |
| import numpy as np | |
| from src.models.segment import Segment | |
| from src.models.difference import Difference, GenerationResult | |
| from src.services.segmentation import SegmentationService | |
| from src.services.saliency import SaliencyService | |
| from src.services.inpainting import InpaintingService | |
| from src.services.color_changer import ColorChanger | |
| from src.services.object_duplicator import ObjectDuplicator | |
| from src.services.quality_evaluator import QualityEvaluator | |
| logger = logging.getLogger(__name__) | |
| ProgressCallback = Callable[[int, str], None] | None | |
| class DifferenceGenerator: | |
| """Generates spot-the-difference images from a single input image.""" | |
| def __init__( | |
| self, | |
| segmentation: SegmentationService, | |
| saliency: SaliencyService, | |
| inpainting: InpaintingService, | |
| color_changer: ColorChanger, | |
| object_duplicator: ObjectDuplicator, | |
| difficulty_config: dict, | |
| segment_min_area_ratio: float = 0.002, | |
| segment_max_area_ratio: float = 0.15, | |
| ) -> None: | |
| self._seg = segmentation | |
| self._sal = saliency | |
| self._inp = inpainting | |
| self._col = color_changer | |
| self._dup = object_duplicator | |
| self._quality = QualityEvaluator() | |
| self._difficulty_config = difficulty_config | |
| self._min_area_ratio = segment_min_area_ratio | |
| self._max_area_ratio = segment_max_area_ratio | |
| def generate( | |
| self, | |
| image: np.ndarray, | |
| difficulty: str = "medium", | |
| progress: ProgressCallback = None, | |
| ) -> GenerationResult: | |
| """Run the full generation pipeline. | |
| Args: | |
| image: BGR image (H, W, 3) uint8. | |
| difficulty: "easy", "medium", or "hard". | |
| progress: Optional callback(percent, step_name). | |
| Returns: | |
| GenerationResult with original, modified image, and difference metadata. | |
| """ | |
| timings: dict[str, float] = {} | |
| _notify(progress, 5, "セグメンテーション開始...") | |
| # 1. Segmentation | |
| t0 = time.time() | |
| segments = self._seg.segment( | |
| image, | |
| min_area_ratio=self._min_area_ratio, | |
| max_area_ratio=self._max_area_ratio, | |
| ) | |
| timings["segmentation"] = time.time() - t0 | |
| _notify(progress, 40, f"セグメンテーション完了 ({len(segments)}個検出)") | |
| if len(segments) < 2: | |
| logger.warning("Too few segments (%d) to generate differences.", len(segments)) | |
| return GenerationResult( | |
| original_image=image, | |
| modified_image=image.copy(), | |
| differences=[], | |
| metadata={"error": "検出されたオブジェクトが少なすぎます"}, | |
| ) | |
| # 2. Saliency analysis | |
| t0 = time.time() | |
| saliency_map = self._sal.compute_map(image) | |
| ranked = self._sal.rank_segments(segments, saliency_map) | |
| timings["saliency"] = time.time() - t0 | |
| _notify(progress, 50, "顕著性解析完了") | |
| # 3. Select segments based on difficulty | |
| selected = self._select_segments(ranked, difficulty) | |
| _notify(progress, 55, f"{len(selected)}個のオブジェクトを変更します") | |
| # 4. Apply changes | |
| t0 = time.time() | |
| modified, differences = self._apply_changes(image, selected, progress) | |
| timings["changes"] = time.time() - t0 | |
| total_time = sum(timings.values()) | |
| timings["total"] = total_time | |
| _notify(progress, 95, "メタデータを生成中...") | |
| metadata = { | |
| "difficulty": difficulty, | |
| "processing_times": {k: round(v, 2) for k, v in timings.items()}, | |
| "segments_detected": len(segments), | |
| "model_versions": { | |
| "segmentation": "FastSAM-x", | |
| "saliency": "OpenCV SpectralResidual", | |
| "inpainting": "OpenCV Navier-Stokes", | |
| }, | |
| } | |
| _notify(progress, 100, "完了") | |
| return GenerationResult( | |
| original_image=image, | |
| modified_image=modified, | |
| differences=differences, | |
| metadata=metadata, | |
| ) | |
| def _select_segments( | |
| self, ranked: list[Segment], difficulty: str | |
| ) -> list[Segment]: | |
| """Pick segments based on difficulty settings with quality filtering.""" | |
| config = self._difficulty_config[difficulty] | |
| num_changes = config["num_changes"] | |
| max_saliency = config["max_saliency"] | |
| # First, filter by quality | |
| quality_filtered = [] | |
| for seg in ranked: | |
| is_acceptable, quality_score, reason = self._quality.evaluate_segment_quality( | |
| np.zeros((seg.mask.shape[0], seg.mask.shape[1], 3), dtype=np.uint8), | |
| seg, | |
| ) | |
| if is_acceptable: | |
| seg.quality_score = quality_score # Store for reference | |
| quality_filtered.append(seg) | |
| else: | |
| logger.debug(f"Segment {seg.id} rejected: {reason}") | |
| logger.info(f"Quality filtered: {len(quality_filtered)} / {len(ranked)} segments passed") | |
| # Then filter by saliency | |
| candidates = [s for s in quality_filtered if s.saliency_score <= max_saliency] | |
| # If not enough candidates below threshold, relax and take from full quality list | |
| if len(candidates) < num_changes: | |
| candidates = quality_filtered.copy() | |
| # If still not enough, take from ranked list (without quality filter) | |
| if len(candidates) < num_changes: | |
| logger.warning(f"Not enough high-quality segments, relaxing quality requirements") | |
| candidates = ranked.copy() | |
| n = min(num_changes, len(candidates)) | |
| return random.sample(candidates, n) | |
| def _apply_changes( | |
| self, | |
| image: np.ndarray, | |
| segments: list[Segment], | |
| progress: ProgressCallback, | |
| ) -> tuple[np.ndarray, list[Difference]]: | |
| """Apply a random change to each selected segment with quality checking.""" | |
| modified = image.copy() | |
| differences: list[Difference] = [] | |
| total = len(segments) | |
| successful_changes = 0 | |
| max_retries = 2 | |
| for i, seg in enumerate(segments): | |
| pct = 55 + int((i / max(total, 1)) * 35) | |
| change_type = self._decide_change_type(seg) | |
| _notify(progress, pct, f"変更を適用中 ({i + 1}/{total}): {change_type}") | |
| # Try to apply change with quality check | |
| for attempt in range(max_retries): | |
| temp_modified = modified.copy() | |
| x1, y1, x2, y2 = seg.bbox | |
| original_region = image[y1:y2, x1:x2].copy() | |
| new_bbox_result = None | |
| # Apply the change | |
| if change_type == "deletion": | |
| temp_modified = self._inp.inpaint(temp_modified, seg.mask) | |
| elif change_type == "color_change": | |
| temp_modified, _ = self._col.change_hue(temp_modified, seg.mask) | |
| elif change_type == "addition": | |
| temp_modified, new_bbox_result = self._dup.duplicate(temp_modified, seg) | |
| if new_bbox_result is None: | |
| logger.debug(f"Addition failed for segment {seg.id}, trying different type") | |
| change_type = random.choice(["deletion", "color_change"]) | |
| continue | |
| # Check quality of the modification | |
| modified_region = temp_modified[y1:y2, x1:x2] | |
| local_mask = seg.mask[y1:y2, x1:x2] | |
| is_acceptable, quality_score, reason = self._quality.evaluate_modification_quality( | |
| original_region, | |
| modified_region, | |
| local_mask, | |
| change_type, | |
| ) | |
| if is_acceptable or attempt == max_retries - 1: | |
| # Accept this modification | |
| modified = temp_modified | |
| diff = self._apply_single_change(modified, seg, change_type, successful_changes + 1) | |
| if diff is not None: | |
| if change_type == "addition" and new_bbox_result is not None: | |
| diff.bbox = new_bbox_result | |
| differences.append(diff) | |
| successful_changes += 1 | |
| if not is_acceptable: | |
| logger.warning(f"Accepting lower quality change ({reason}) after {max_retries} attempts") | |
| else: | |
| logger.debug(f"Change accepted with quality score: {quality_score:.2f}") | |
| break | |
| else: | |
| logger.debug(f"Modification rejected ({reason}), retrying with different parameters...") | |
| # Try a different change type on retry | |
| if change_type == "color_change": | |
| change_type = random.choice(["deletion", "addition"]) | |
| elif change_type == "addition": | |
| change_type = "color_change" | |
| logger.info(f"Successfully applied {successful_changes} / {total} changes") | |
| return modified, differences | |
| def _apply_single_change( | |
| self, | |
| image: np.ndarray, | |
| seg: Segment, | |
| change_type: str, | |
| diff_id: int, | |
| ) -> Difference | None: | |
| """Create a Difference record. Actual mutation happens in _apply_changes.""" | |
| descriptions = { | |
| "deletion": "オブジェクトを削除", | |
| "color_change": "色を変更", | |
| "addition": "オブジェクトを追加", | |
| } | |
| return Difference( | |
| id=diff_id, | |
| type=change_type, | |
| bbox=seg.bbox, | |
| saliency_score=seg.saliency_score, | |
| description=descriptions.get(change_type, ""), | |
| ) | |
| def _decide_change_type(self, seg: Segment) -> str: | |
| """Choose change type, preferring colour/addition for large objects.""" | |
| image_area_ratio = seg.area / max(seg.mask.size, 1) | |
| if image_area_ratio > 0.08: | |
| # Large object — inpainting quality degrades, avoid deletion | |
| return random.choice(["color_change", "addition"]) | |
| return random.choice(["deletion", "color_change", "addition"]) | |
| def _notify(cb: ProgressCallback, percent: int, step: str) -> None: | |
| if cb is not None: | |
| cb(percent, step) | |