""" pipeline.py ------------ Core orchestrator -- ties together detection, segmentation, and inpainting. Usage (programmatic): from pipeline import ObjectRemovalPipeline pipe = ObjectRemovalPipeline() result = pipe.run(scene_path, object_paths) result.save("output.png") Usage (CLI): python pipeline.py --scene input/scene/room.jpg \ --objects input/objects/chair.jpg input/objects/lamp.jpg """ import os import sys import argparse from pathlib import Path from typing import List, Optional, Tuple import numpy as np from PIL import Image from config import ( OUTPUT_DIR, INPAINT_METHOD, SAVE_DEBUG_IMAGES, CLIP_SIMILARITY_THRESHOLD, MASK_DILATION_PX, ) from clip_matcher import CLIPMatcher from detector import GroundingDINODetector from segmenter import SAMSegmenter from image_utils import ( load_image_pil, save_image, boxes_to_mask, combine_masks, save_detection_debug, save_mask_debug, save_comparison, list_images, ) # -- Inpainter factory --------------------------------------------------------- def _make_inpainter(): if INPAINT_METHOD == "lama": from inpainter_lama import LamaInpainter return LamaInpainter() elif INPAINT_METHOD == "sd_inpaint": from inpainter_sd import SDInpainter return SDInpainter() else: raise ValueError(f"Unknown INPAINT_METHOD: {INPAINT_METHOD!r}") # -- Pipeline ------------------------------------------------------------------ class ObjectRemovalPipeline: """ End-to-end object removal: reference images CLIP label GroundingDINO detection CLIP similarity filter SAM segmentation inpainting """ def __init__(self) -> None: self.clip = CLIPMatcher() self.detector = GroundingDINODetector() self.segmenter = SAMSegmenter() self.inpainter = _make_inpainter() # -- Public entry point ---------------------------------------------------- def run( self, scene_path: str, object_paths: List[str], output_dir: str = OUTPUT_DIR, save_debug: bool = SAVE_DEBUG_IMAGES, threshold: Optional[float] = None, ) -> Image.Image: """ Remove all objects specified by `object_paths` from `scene_path`. Returns the inpainted PIL image and writes it to `output_dir`. """ os.makedirs(output_dir, exist_ok=True) stem = Path(scene_path).stem print(f"\n{'='*60}") print(f" Scene : {scene_path}") print(f" Objects ({len(object_paths)}): {[os.path.basename(p) for p in object_paths]}") print(f"{'='*60}") scene_pil = load_image_pil(scene_path) w, h = scene_pil.size # -- Stage 1: detect & match each reference object ----------------- all_boxes = [] all_dets = [] # for debug visualisation for obj_path in object_paths: print(f"\n-- Object: {os.path.basename(obj_path)}") obj_pil = load_image_pil(obj_path) # 1a + 1b. Detect in scene using Reference Image Query (Image-Guided Detection) dets = self.detector.detect_from_image(scene_pil, obj_pil) candidate_boxes = [d["box"] for d in dets] all_dets.extend(dets) if not candidate_boxes: print(f" [!] No candidates found for image query -- skipping") continue # 1c. Filter by CLIP image-image similarity target_threshold = threshold if threshold is not None else CLIP_SIMILARITY_THRESHOLD accepted = self.clip.filter_boxes_by_similarity( obj_pil, scene_pil, candidate_boxes, threshold=target_threshold ) if not accepted: print(f" [!] No boxes passed similarity threshold -- skipping") else: print(f" [v] {len(accepted)} box(es) accepted") all_boxes.extend(accepted) if not all_boxes: print("\n[!] No objects detected. Returning original image.") return scene_pil # -- Debug: save detection visualisation --------------------------- if save_debug: det_path = os.path.join(output_dir, f"{stem}_debug_detections.jpg") save_detection_debug(scene_path, all_dets, det_path) print(f"\n [debug] detection image {det_path}") # -- Stage 2: SAM segmentation ------------------------------------- print(f"\n-- Segmentation ({len(all_boxes)} box(es))") combined_mask = self.segmenter.segment_boxes(scene_pil, all_boxes) if save_debug: mask_path = os.path.join(output_dir, f"{stem}_debug_mask.jpg") save_mask_debug(scene_path, combined_mask, mask_path) print(f" [debug] mask image {mask_path}") # -- Stage 3: Inpainting ------------------------------------------- print(f"\n-- Inpainting ({INPAINT_METHOD})") result_pil = self.inpainter.inpaint(scene_pil, combined_mask) # -- Save outputs -------------------------------------------------- out_path = os.path.join(output_dir, f"{stem}_result.png") result_pil.save(out_path) print(f"\n [DONE] Result saved {out_path}") if save_debug: cmp_path = os.path.join(output_dir, f"{stem}_comparison.jpg") save_comparison(scene_pil, result_pil, cmp_path, labels=("Original", "Objects Removed")) print(f" [debug] comparison {cmp_path}") return result_pil # -- Batch helper ---------------------------------------------------------- def run_batch( self, scene_dir: str, objects_dir: str, output_dir: str = OUTPUT_DIR, ) -> None: """ Process all images in `scene_dir`, removing all objects found in `objects_dir`. """ scenes = list_images(scene_dir) objects = list_images(objects_dir) if not scenes: print(f"No images found in scene directory: {scene_dir}") return if not objects: print(f"No object images found in: {objects_dir}") return print(f"Batch: {len(scenes)} scene(s) {len(objects)} object reference(s)") for scene_path in scenes: self.run(scene_path, objects, output_dir=output_dir) # -- CLI ----------------------------------------------------------------------- def _parse_args(): parser = argparse.ArgumentParser( description="Object Removal Pipeline -- remove specific objects from a scene image." ) group = parser.add_mutually_exclusive_group(required=True) group.add_argument( "--scene", type=str, help="Path to the scene image." ) group.add_argument( "--scene-dir", type=str, help="Directory of scene images (batch mode)." ) parser.add_argument( "--objects", type=str, nargs="+", help="One or more paths to reference object images." ) parser.add_argument( "--objects-dir", type=str, help="Directory of object images (alternative to --objects)." ) parser.add_argument( "--output-dir", type=str, default=OUTPUT_DIR, help=f"Output directory (default: {OUTPUT_DIR})." ) parser.add_argument( "--inpaint", choices=["lama", "sd_inpaint"], default=None, help="Override inpainting method from config." ) parser.add_argument( "--no-debug", action="store_true", help="Skip saving debug images." ) return parser.parse_args() def main(): args = _parse_args() # Override config if flags given if args.inpaint: import config config.INPAINT_METHOD = args.inpaint pipe = ObjectRemovalPipeline() # Resolve object paths if args.objects: object_paths = args.objects elif args.objects_dir: object_paths = list_images(args.objects_dir) else: print("Error: provide --objects or --objects-dir") sys.exit(1) save_debug = not args.no_debug if args.scene: pipe.run(args.scene, object_paths, output_dir=args.output_dir, save_debug=save_debug) elif args.scene_dir: pipe.run_batch(args.scene_dir, args.objects_dir or "", output_dir=args.output_dir) if __name__ == "__main__": main()