import argparse import csv import sys from pathlib import Path import numpy as np from PIL import Image sys.path.insert(0, str(Path(__file__).resolve().parents[1])) from common.image_io import list_images, read_rgb from crop.rfdetr_adapter import detections_to_records, load_rfdetr_seg_model, predict_records from crop.utils import center_fallback_bbox def mask_bbox(mask: np.ndarray) -> tuple[int, int, int, int] | None: ys, xs = np.where(mask) if len(xs) == 0 or len(ys) == 0: return None return int(xs.min()), int(ys.min()), int(xs.max()) + 1, int(ys.max()) + 1 def resize_mask(mask: np.ndarray, image_size: tuple[int, int]) -> np.ndarray: if mask.shape == (image_size[1], image_size[0]): return mask.astype(bool) mask_image = Image.fromarray(mask.astype("uint8") * 255) return np.asarray(mask_image.resize(image_size, Image.Resampling.NEAREST)) > 0 def save_masked_outputs( image: Image.Image, mask: np.ndarray, masked_path: Path, crop_path: Path, alpha_path: Path | None = None, ) -> tuple[int, int, int, int]: mask = resize_mask(mask, image.size) bbox = mask_bbox(mask) if bbox is None: raise RuntimeError("empty soup_valid_area mask") rgb = np.asarray(image) masked = np.zeros_like(rgb) masked[mask] = rgb[mask] masked_path.parent.mkdir(parents=True, exist_ok=True) Image.fromarray(masked).save(masked_path) Image.fromarray(masked).crop(bbox).save(crop_path) if alpha_path: rgba = np.dstack([rgb, mask.astype("uint8") * 255]) Image.fromarray(rgba).save(alpha_path) return bbox def extract_one_image( image_path: str | Path, output_dir: str | Path, checkpoint_path: str | None = None, model_size: str = "small", threshold: float = 0.5, low_confidence_threshold: float = 0.35, predictor=None, save_alpha: bool = True, preserve_filename: bool = False, output_kind: str = "all", skip_fallback_output: bool = False, ) -> dict: image_path = Path(image_path) output_dir = Path(output_dir) image = read_rgb(image_path) if preserve_filename: primary_path = output_dir / image_path.name masked_path = primary_path if output_kind in {"crop", "masked"} else output_dir / f"{image_path.stem}_soup_masked.jpg" crop_path = primary_path if output_kind == "crop" else output_dir / f"{image_path.stem}_soup_crop.jpg" alpha_path = primary_path.with_suffix(".png") if output_kind == "alpha" else None else: masked_path = output_dir / f"{image_path.stem}_soup_masked.jpg" crop_path = output_dir / f"{image_path.stem}_soup_crop.jpg" alpha_path = output_dir / f"{image_path.stem}_soup_alpha.png" if save_alpha else None try: records = predictor(image_path) if predictor else predict_records(image_path, model_size, checkpoint_path, threshold) records = [rec for rec in records if rec.get("mask") is not None] if not records: raise RuntimeError("no soup_valid_area mask detection") best = max(records, key=lambda item: item.get("confidence", 0.0)) bbox = save_masked_outputs(image, np.asarray(best["mask"]).astype(bool), masked_path, crop_path, alpha_path) if preserve_filename and output_kind == "masked": Path(crop_path).unlink(missing_ok=True) confidence = float(best.get("confidence", 0.0)) status = 1 if confidence >= low_confidence_threshold else 2 error_message = "" except Exception as exc: bbox = center_fallback_bbox(image.size) if not skip_fallback_output: masked_path.parent.mkdir(parents=True, exist_ok=True) fallback_output = crop_path if preserve_filename and output_kind == "crop" else masked_path if preserve_filename and output_kind == "masked": image.save(fallback_output) else: image.crop(bbox).save(fallback_output) confidence = 0.0 status = 0 error_message = str(exc) x1, y1, x2, y2 = bbox return { "image_id": image_path.stem, "source_image_path": str(image_path), "soup_masked_path": str(masked_path), "soup_crop_path": str(crop_path), "soup_alpha_path": str(alpha_path) if alpha_path else "", "extract_status": status, "extract_confidence": confidence, "x1": x1, "y1": y1, "x2": x2, "y2": y2, "error_message": error_message, } def build_predictor(model_size: str, checkpoint_path: str, threshold: float): model = load_rfdetr_seg_model(model_size=model_size, checkpoint_path=checkpoint_path) def predictor(image_path: str | Path) -> list[dict]: detections = model.predict(str(image_path), threshold=threshold) return detections_to_records(detections) return predictor def main() -> int: parser = argparse.ArgumentParser(description="Extract soup valid area images with an RF-DETR-Seg mask model.") parser.add_argument("--input", required=True) parser.add_argument("--output-dir", default="data/processed_images/soup_valid_area") parser.add_argument("--report", default="reports/crop/soup_valid_area/extract_result.csv") parser.add_argument("--checkpoint-path", required=True) parser.add_argument("--model-size", default="small", choices=["small", "medium"]) parser.add_argument("--threshold", type=float, default=0.5) parser.add_argument("--low-confidence-threshold", type=float, default=0.35) parser.add_argument("--no-alpha", action="store_true") parser.add_argument("--preserve-filenames", action="store_true") parser.add_argument("--output-kind", default="all", choices=["all", "crop", "masked", "alpha"]) parser.add_argument("--skip-fallback-output", action="store_true") args = parser.parse_args() predictor = build_predictor(args.model_size, args.checkpoint_path, args.threshold) rows = [ extract_one_image( image_path, args.output_dir, args.checkpoint_path, args.model_size, args.threshold, args.low_confidence_threshold, predictor=predictor, save_alpha=not args.no_alpha, preserve_filename=args.preserve_filenames, output_kind=args.output_kind, skip_fallback_output=args.skip_fallback_output, ) for image_path in list_images(args.input) ] fieldnames = [ "image_id", "source_image_path", "soup_masked_path", "soup_crop_path", "soup_alpha_path", "extract_status", "extract_confidence", "x1", "y1", "x2", "y2", "error_message", ] report = Path(args.report) report.parent.mkdir(parents=True, exist_ok=True) with report.open("w", newline="", encoding="utf-8") as f: writer = csv.DictWriter(f, fieldnames=fieldnames) writer.writeheader() writer.writerows(rows) print(f"wrote soup extraction report: {report}") return 0 if __name__ == "__main__": raise SystemExit(main())