"""Run box-prompted segmentation with EfficientSAM-S and ONNX Runtime.""" import argparse import json from pathlib import Path import numpy as np import onnxruntime as ort from PIL import Image, ImageDraw MODEL_DIR = Path(__file__).resolve().parent MODEL_PATH = MODEL_DIR / "efficient_sam_s_graviton_onnx_optimized.onnx" SAMPLE_IMAGE_PATH = MODEL_DIR / "sample_input.jpg" INPUT_SIZE = 1024 SAMPLE_BOX_PROMPT = [410, 510, 600, 830] EXECUTION_PROVIDER = "CPUExecutionProvider" def required_file(path: Path) -> Path: if not path.is_file(): raise FileNotFoundError(f"required model file not found: {path}") return path def load_model() -> ort.InferenceSession: available = ort.get_available_providers() if EXECUTION_PROVIDER not in available: raise RuntimeError( f"{EXECUTION_PROVIDER} is unavailable; available providers: {', '.join(available)}" ) session = ort.InferenceSession(str(required_file(MODEL_PATH)), providers=[EXECUTION_PROVIDER]) input_names = [node.name for node in session.get_inputs()] if input_names != ["image", "point_coords", "point_labels"]: raise RuntimeError(f"unexpected model inputs: {input_names}") return session def resize_bilinear(image: np.ndarray, size: int) -> np.ndarray: """Resize an HWC float image with half-pixel centres and no antialiasing.""" height, width = image.shape[:2] def axis(source: int) -> tuple[np.ndarray, np.ndarray, np.ndarray]: coords = (np.arange(size, dtype=np.float64) + 0.5) * (source / size) - 0.5 coords = np.clip(coords, 0.0, source - 1) low = np.floor(coords).astype(np.int64) high = np.minimum(low + 1, source - 1) return low, high, (coords - low).astype(np.float32) y0, y1, wy = axis(height) x0, x1, wx = axis(width) wy = wy[:, None, None] wx = wx[None, :, None] top = image[y0][:, x0] * (1 - wx) + image[y0][:, x1] * wx bottom = image[y1][:, x0] * (1 - wx) + image[y1][:, x1] * wx return top * (1 - wy) + bottom * wy def preprocess(image_path: Path) -> np.ndarray: with Image.open(required_file(image_path)) as image: pixels = np.asarray(image.convert("RGB"), dtype=np.float32) / 255.0 resized = resize_bilinear(pixels, INPUT_SIZE) return np.ascontiguousarray(resized.transpose(2, 0, 1)[None], dtype=np.float32) def make_box_prompt(box: list[int]) -> tuple[np.ndarray, np.ndarray]: x1, y1, x2, y2 = box point_coords = np.array([[[[x1, y1], [x2, y2]]]], dtype=np.float32) point_labels = np.array([[[2, 3]]], dtype=np.float32) return point_coords, point_labels def validate_box(box: list[int]) -> list[int]: x1, y1, x2, y2 = box if not (0 <= x1 < x2 <= INPUT_SIZE and 0 <= y1 < y2 <= INPUT_SIZE): raise ValueError( f"box must satisfy 0 <= x1 < x2 <= {INPUT_SIZE} and " f"0 <= y1 < y2 <= {INPUT_SIZE}; received {box}" ) return box def postprocess(raw_output: np.ndarray, box: list[int]) -> tuple[np.ndarray, dict[str, object]]: if tuple(raw_output.shape) != (1, INPUT_SIZE, INPUT_SIZE): raise ValueError( f"expected output shape (1, {INPUT_SIZE}, {INPUT_SIZE}), " f"found {tuple(raw_output.shape)}" ) if not np.isfinite(raw_output).all(): raise ValueError("model output contains non-finite mask values") mask = (raw_output[0] > 0.5).astype(np.uint8) pixel_count = int(mask.sum()) total_pixels = int(mask.size) result = { "pixel_count": pixel_count, "total_pixels": total_pixels, "coverage_pct": round(100.0 * pixel_count / total_pixels, 2), "box_prompt": box, } return mask, result def create_overlay(mask: np.ndarray, image_path: Path, box: list[int]) -> Image.Image: with Image.open(required_file(image_path)) as image: source = image.convert("RGB") source_width, source_height = source.size source_mask = Image.fromarray(mask * 255).resize( source.size, resample=Image.Resampling.NEAREST, ) source_mask_array = np.asarray(source_mask, dtype=np.uint8) > 0 mask_rgba = np.zeros((source_height, source_width, 4), dtype=np.uint8) mask_rgba[source_mask_array] = [0, 200, 0, 200] overlay = Image.fromarray(mask_rgba) composited = Image.alpha_composite(source.convert("RGBA"), overlay).convert("RGB") x1, y1, x2, y2 = box scaled_box = [ round(x1 * source_width / INPUT_SIZE), round(y1 * source_height / INPUT_SIZE), round(x2 * source_width / INPUT_SIZE), round(y2 * source_height / INPUT_SIZE), ] draw = ImageDraw.Draw(composited) draw.rectangle(scaled_box, outline=(255, 0, 0), width=3) return composited def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser(description=__doc__) parser.add_argument( "--image", type=Path, default=SAMPLE_IMAGE_PATH, help="input image (default: sample_input.jpg)", ) parser.add_argument( "--box", nargs=4, type=int, metavar=("X1", "Y1", "X2", "Y2"), default=SAMPLE_BOX_PROMPT, help="box prompt in resized 1024x1024 coordinates (default: sample box)", ) parser.add_argument("--output-dir", type=Path, help="optional output directory") return parser.parse_args() def main() -> None: args = parse_args() image_path = required_file(args.image) box = validate_box(args.box) image_tensor = preprocess(image_path) point_coords, point_labels = make_box_prompt(box) session = load_model() raw_output = session.run( None, {"image": image_tensor, "point_coords": point_coords, "point_labels": point_labels}, )[0] mask, result = postprocess(raw_output, box) print(f"Input: {image_path.name}; tensor shape: {tuple(image_tensor.shape)}") print(f"Box prompt: {box}") print( f"Segmented pixels: {result['pixel_count']:,} / {result['total_pixels']:,} " f"({result['coverage_pct']:.2f}%)" ) if args.output_dir is not None: args.output_dir.mkdir(parents=True, exist_ok=True) overlay_path = args.output_dir / "sample_output.png" result_path = args.output_dir / "segmentation.json" create_overlay(mask, image_path, box).save(overlay_path) result_path.write_text(json.dumps(result, indent=2) + "\n", encoding="utf-8") print(f"Saved: {overlay_path}") print(f"Saved: {result_path}") if __name__ == "__main__": main()