| """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() |
|
|
