| |
| """Clean, merge, and regularise window detections from a segmentation model. |
| |
| Windows are the most safety-critical keep-out for a painting robot, and the |
| semantic model gets them wrong in predictable ways: a window split by mullions |
| comes out as several fragments, thin frames erode or break, and diagonal or |
| perspective views never produce a clean rectangle. Raw connected components |
| inherit all of those flaws. |
| |
| This module fixes the output geometry without retraining: |
| |
| * ``merge_fragments`` joins pieces separated by up to a gap (mullion bars, |
| occlusions) into one window; |
| * ``minimum_area_rectangle`` fits an oriented rectangle with rotating calipers, |
| so a window stays a window at any viewing angle; |
| * ``window_instances`` runs clean -> merge -> fit -> filter and returns |
| instances with both axis-aligned boxes and oriented quads; |
| * ``regularized_mask`` rasterises the fitted rectangles into a clean keep-out |
| mask. |
| |
| numpy + scipy only, so it is unit-testable and runs anywhere. See |
| ``python3 window_postprocess.py --help`` for the CLI. |
| """ |
| from __future__ import annotations |
|
|
| import argparse |
| import json |
| from pathlib import Path |
|
|
| import numpy as np |
| from PIL import Image, ImageDraw |
| from scipy import ndimage |
|
|
|
|
| def disk(radius): |
| size = int(radius) * 2 + 1 |
| ys, xs = np.mgrid[0:size, 0:size] |
| return (xs - radius) ** 2 + (ys - radius) ** 2 <= radius ** 2 |
|
|
|
|
| def clean_mask(mask, min_area=0, open_radius=0, close_radius=0): |
| """Morphological tidy-up: close gaps, open specks, drop small components.""" |
| mask = np.asarray(mask, dtype=bool) |
| if close_radius > 0: |
| mask = ndimage.binary_closing(mask, structure=disk(close_radius)) |
| if open_radius > 0: |
| mask = ndimage.binary_opening(mask, structure=disk(open_radius)) |
| if min_area > 0: |
| labels, count = ndimage.label(mask, structure=np.ones((3, 3), dtype=int)) |
| if count: |
| sizes = ndimage.sum(mask, labels, index=np.arange(1, count + 1)) |
| keep = np.flatnonzero(sizes >= min_area) + 1 |
| mask = np.isin(labels, keep) |
| return mask |
|
|
|
|
| def merge_fragments(mask, gap): |
| """Merge components separated by up to ``gap`` pixels (e.g. mullion bars).""" |
| mask = np.asarray(mask, dtype=bool) |
| if gap <= 0 or not mask.any(): |
| labels, _ = ndimage.label(mask, structure=np.ones((3, 3), dtype=int)) |
| return labels |
| radius = max(1, int(round(gap / 2.0))) |
| dilated = ndimage.binary_dilation(mask, structure=disk(radius)) |
| merged, _ = ndimage.label(dilated, structure=np.ones((3, 3), dtype=int)) |
| labels = np.where(mask, merged, 0) |
| return labels |
|
|
|
|
| def convex_hull(points): |
| """Monotone-chain hull of (N, 2) points, returned counter-clockwise.""" |
| unique = sorted({(float(x), float(y)) for x, y in points}) |
| if len(unique) <= 2: |
| return np.asarray(unique, float) |
|
|
| def cross(o, a, b): |
| return (a[0] - o[0]) * (b[1] - o[1]) - (a[1] - o[1]) * (b[0] - o[0]) |
|
|
| lower = [] |
| for point in unique: |
| while len(lower) >= 2 and cross(lower[-2], lower[-1], point) <= 0: |
| lower.pop() |
| lower.append(point) |
| upper = [] |
| for point in reversed(unique): |
| while len(upper) >= 2 and cross(upper[-2], upper[-1], point) <= 0: |
| upper.pop() |
| upper.append(point) |
| return np.asarray(lower[:-1] + upper[:-1], float) |
|
|
|
|
| def minimum_area_rectangle(points): |
| """Oriented rectangle of least area around points (rotating calipers). |
| |
| Returns ``(corners, width, height, angle_deg)`` where ``corners`` is a |
| (4, 2) array. Robust to any orientation, unlike an axis-aligned box. |
| """ |
| points = np.asarray(points, float) |
| hull = convex_hull(points) |
| if len(hull) < 3: |
| x0, y0 = points.min(axis=0) |
| x1, y1 = points.max(axis=0) |
| corners = np.array([[x0, y0], [x1, y0], [x1, y1], [x0, y1]], float) |
| return corners, float(x1 - x0), float(y1 - y0), 0.0 |
| best = None |
| count = len(hull) |
| for i in range(count): |
| p1, p2 = hull[i], hull[(i + 1) % count] |
| edge = p2 - p1 |
| length = float(np.hypot(*edge)) |
| if length == 0: |
| continue |
| axis_x = edge / length |
| axis_y = np.array([-axis_x[1], axis_x[0]]) |
| relative = hull - p1 |
| proj_x = relative @ axis_x |
| proj_y = relative @ axis_y |
| width = float(proj_x.max() - proj_x.min()) |
| height = float(proj_y.max() - proj_y.min()) |
| area = width * height |
| if best is None or area < best[0]: |
| center = p1 + (0.5 * (proj_x.min() + proj_x.max())) * axis_x \ |
| + (0.5 * (proj_y.min() + proj_y.max())) * axis_y |
| half_w, half_h = width / 2.0, height / 2.0 |
| corners = np.array([center - half_w * axis_x - half_h * axis_y, |
| center + half_w * axis_x - half_h * axis_y, |
| center + half_w * axis_x + half_h * axis_y, |
| center - half_w * axis_x + half_h * axis_y]) |
| best = (area, corners, width, height, float(np.degrees(np.arctan2(axis_x[1], axis_x[0])))) |
| return best[1], best[2], best[3], best[4] |
|
|
|
|
| def window_instances(mask, confidence=None, min_area=64, merge_gap=8, open_radius=1, close_radius=2, |
| min_fill=0.5, max_aspect=8.0, max_instances=200): |
| """Detect clean window instances from a binary mask (and optional probability). |
| |
| Pipeline: morphological clean -> fragment merge -> oriented rectangle fit -> |
| filter by area, fill ratio (a window fills its own rectangle), and aspect |
| ratio. Returns ``(instances, labels)`` where ``labels`` is a per-pixel map |
| whose positive values join each kept instance. |
| """ |
| mask = clean_mask(np.asarray(mask, dtype=bool), min_area=0, |
| open_radius=open_radius, close_radius=close_radius) |
| labels = merge_fragments(mask, merge_gap) |
| instances = [] |
| for label in range(1, int(labels.max()) + 1 if labels.size else 0): |
| component = labels == label |
| area = int(component.sum()) |
| if area < min_area: |
| continue |
| ys, xs = np.nonzero(component) |
| corners, width, height, angle = minimum_area_rectangle(np.c_[xs, ys]) |
| if width <= 0 or height <= 0: |
| continue |
| fill = area / (width * height) |
| aspect = max(width, height) / max(1e-6, min(width, height)) |
| if fill < min_fill or aspect > max_aspect: |
| continue |
| score = float(confidence[component].mean()) if confidence is not None else None |
| instances.append({ |
| "label": label, |
| "bbox_xyxy": [int(xs.min()), int(ys.min()), int(xs.max()), int(ys.max())], |
| "quad_xy": [[round(float(x), 2), round(float(y), 2)] for x, y in corners], |
| "center_xy": [round(float(xs.mean()), 2), round(float(ys.mean()), 2)], |
| "size_xy": [round(float(width), 2), round(float(height), 2)], |
| "angle_deg": round(float(angle), 2), |
| "fill_ratio": round(float(fill), 4), |
| "area_pixels": area, |
| "mean_confidence": round(score, 4) if score is not None else None, |
| }) |
| instances.sort(key=lambda item: item["area_pixels"], reverse=True) |
| instances = instances[:max_instances] |
| kept = {item["label"] for item in instances} |
| kept_labels = np.where(np.isin(labels, list(kept) or [0]), labels, 0) |
| return instances, kept_labels |
|
|
|
|
| def regularized_mask(instances, shape, use_quad=True): |
| """Rasterise fitted rectangles into a clean binary window mask.""" |
| canvas = Image.new("L", (shape[1], shape[0]), 0) |
| draw = ImageDraw.Draw(canvas) |
| for instance in instances: |
| if use_quad: |
| draw.polygon([tuple(point) for point in instance["quad_xy"]], fill=255) |
| else: |
| x0, y0, x1, y1 = instance["bbox_xyxy"] |
| draw.rectangle([x0, y0, x1, y1], fill=255) |
| return np.asarray(canvas) > 0 |
|
|
|
|
| def window_metrics(prediction, target): |
| """Pixel precision/recall/IoU plus instance counts for window evaluation.""" |
| prediction, target = np.asarray(prediction, bool), np.asarray(target, bool) |
| true_positive = int((prediction & target).sum()) |
| false_positive = int((prediction & ~target).sum()) |
| false_negative = int((~prediction & target).sum()) |
| precision = true_positive / (true_positive + false_positive) if true_positive + false_positive else 1.0 |
| recall = true_positive / (true_positive + false_negative) if true_positive + false_negative else 1.0 |
| iou = true_positive / (true_positive + false_positive + false_negative) if (true_positive + false_positive + false_negative) else 1.0 |
| return {"precision": round(precision, 4), "recall": round(recall, 4), "iou": round(iou, 4), |
| "predicted_pixels": int(prediction.sum()), "target_pixels": int(target.sum())} |
|
|
|
|
| def main(): |
| parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) |
| parser.add_argument("window", type=Path, help="window probability (grayscale) or binary mask PNG") |
| parser.add_argument("--threshold", type=float, default=0.5, help="probability threshold for a prob map") |
| parser.add_argument("--min-area", type=int, default=64) |
| parser.add_argument("--merge-gap", type=int, default=8, help="merge fragments separated by up to this many pixels") |
| parser.add_argument("--open-radius", type=int, default=1) |
| parser.add_argument("--close-radius", type=int, default=2) |
| parser.add_argument("--min-fill", type=float, default=0.5) |
| parser.add_argument("--max-aspect", type=float, default=8.0) |
| parser.add_argument("--max-instances", type=int, default=200) |
| parser.add_argument("--out-mask", type=Path, help="write the regularized window mask") |
| parser.add_argument("--overlay", type=Path, help="write an overlay of detected windows") |
| parser.add_argument("--json", type=Path, help="write instance JSON") |
| args = parser.parse_args() |
|
|
| raw = np.asarray(Image.open(args.window).convert("L")) |
| probability = raw.astype(np.float32) / 255.0 if raw.max() > 1 else raw.astype(np.float32) |
| mask = probability >= args.threshold |
| instances, labels = window_instances(mask, confidence=probability, min_area=args.min_area, |
| merge_gap=args.merge_gap, open_radius=args.open_radius, |
| close_radius=args.close_radius, min_fill=args.min_fill, |
| max_aspect=args.max_aspect, max_instances=args.max_instances) |
| result = {"windows": len(instances), "instances": instances, |
| "raw_pixels": int(mask.sum()), |
| "regularized_pixels": int(regularized_mask(instances, mask.shape).sum())} |
| rendered = json.dumps(result, indent=2) |
| print(rendered) |
| if args.json: |
| args.json.parent.mkdir(parents=True, exist_ok=True) |
| args.json.write_text(rendered + "\n", encoding="utf-8") |
| if args.out_mask: |
| args.out_mask.parent.mkdir(parents=True, exist_ok=True) |
| Image.fromarray(regularized_mask(instances, mask.shape).astype(np.uint8) * 255).save(args.out_mask) |
| if args.overlay: |
| canvas = np.stack([raw] * 3, axis=-1).astype(np.uint8) |
| canvas[mask] = (canvas[mask] * 0.4 + np.array([0, 0, 160])).astype(np.uint8) |
| image = Image.fromarray(canvas) |
| draw = ImageDraw.Draw(image) |
| for instance in instances: |
| draw.polygon([tuple(point) for point in instance["quad_xy"]], outline=(0, 255, 255)) |
| args.overlay.parent.mkdir(parents=True, exist_ok=True) |
| image.save(args.overlay) |
| print(f"overlay: {args.overlay}") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
