| """Standalone inference API and CLI for StrandSeg-Lite.""" |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| from dataclasses import dataclass |
| from pathlib import Path |
|
|
| import cv2 |
| import numpy as np |
| import torch |
|
|
| from model import load_model |
|
|
|
|
| IMAGE_MEAN = np.asarray([0.485, 0.456, 0.406], dtype=np.float32) |
| IMAGE_STD = np.asarray([0.229, 0.224, 0.225], dtype=np.float32) |
|
|
|
|
| @dataclass(frozen=True) |
| class Prediction: |
| probability: np.ndarray |
| mask: np.ndarray |
| threshold: float |
|
|
|
|
| def _letterbox(image: np.ndarray, width: int, height: int) -> tuple[np.ndarray, tuple[int, int, int, int]]: |
| source_height, source_width = image.shape[:2] |
| scale = min(width / source_width, height / source_height) |
| resized_width = max(1, int(round(source_width * scale))) |
| resized_height = max(1, int(round(source_height * scale))) |
| pad_x = (width - resized_width) // 2 |
| pad_y = (height - resized_height) // 2 |
| resized = cv2.resize(image, (resized_width, resized_height), interpolation=cv2.INTER_AREA) |
| canvas = np.zeros((height, width, 3), dtype=np.uint8) |
| canvas[pad_y : pad_y + resized_height, pad_x : pad_x + resized_width] = resized |
| return canvas, (pad_x, pad_y, resized_width, resized_height) |
|
|
|
|
| class StrandSegLite: |
| def __init__(self, checkpoint: str | Path = "model.pt", device: str | torch.device = "cpu"): |
| self.device = torch.device(device) |
| self.model, self.metadata = load_model(checkpoint, self.device) |
| config = dict(self.metadata.get("model_config") or {}) |
| self.input_width = int(config.get("input_width", 960)) |
| self.input_height = int(config.get("input_height", 544)) |
| self.threshold = float(self.metadata.get("threshold", 0.9204)) |
|
|
| @torch.inference_mode() |
| def predict(self, image: np.ndarray, *, color: str = "rgb", threshold: float | None = None) -> Prediction: |
| if image.ndim != 3 or image.shape[2] != 3: |
| raise ValueError("image must be an HxWx3 uint8 array") |
| if image.dtype != np.uint8: |
| raise ValueError("image must have dtype uint8") |
| color = color.lower() |
| if color == "rgb": |
| bgr = cv2.cvtColor(image, cv2.COLOR_RGB2BGR) |
| elif color == "bgr": |
| bgr = image |
| else: |
| raise ValueError("color must be 'rgb' or 'bgr'") |
| original_height, original_width = bgr.shape[:2] |
| prepared, (pad_x, pad_y, resized_width, resized_height) = _letterbox( |
| bgr, self.input_width, self.input_height |
| ) |
| rgb = cv2.cvtColor(prepared, cv2.COLOR_BGR2RGB).astype(np.float32) / 255.0 |
| rgb = (rgb - IMAGE_MEAN) / IMAGE_STD |
| tensor = torch.from_numpy(np.transpose(rgb, (2, 0, 1))).unsqueeze(0).to(self.device) |
| probability = torch.sigmoid(self.model(tensor))[0, 0].cpu().numpy() |
| probability = probability[pad_y : pad_y + resized_height, pad_x : pad_x + resized_width] |
| probability = cv2.resize(probability, (original_width, original_height), interpolation=cv2.INTER_LINEAR) |
| selected_threshold = self.threshold if threshold is None else float(threshold) |
| mask = (probability >= selected_threshold).astype(np.uint8) |
| return Prediction(probability=probability.astype(np.float32), mask=mask, threshold=selected_threshold) |
|
|
| def predict_file(self, path: str | Path, *, threshold: float | None = None) -> Prediction: |
| encoded = np.fromfile(str(path), dtype=np.uint8) |
| image = cv2.imdecode(encoded, cv2.IMREAD_COLOR) if encoded.size else None |
| if image is None: |
| raise ValueError(f"Could not read image: {path}") |
| return self.predict(image, color="bgr", threshold=threshold) |
|
|
|
|
| def main() -> None: |
| parser = argparse.ArgumentParser(description="Run StrandSeg-Lite on one image") |
| parser.add_argument("image") |
| parser.add_argument("--checkpoint", default="model.pt") |
| parser.add_argument("--device", default="cpu") |
| parser.add_argument("--threshold", type=float) |
| parser.add_argument("--output-mask", default="mask.png") |
| parser.add_argument("--output-probability", default="probability.png") |
| args = parser.parse_args() |
| result = StrandSegLite(args.checkpoint, args.device).predict_file(args.image, threshold=args.threshold) |
| cv2.imwrite(args.output_mask, result.mask * 255) |
| cv2.imwrite(args.output_probability, np.clip(result.probability * 255.0, 0, 255).astype(np.uint8)) |
| print(f"threshold={result.threshold:.4f} foreground_pixels={int(result.mask.sum())}") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
