segment_wall.py

#!/usr/bin/env python3
"""HSV baseline for painted-area estimation on a user-marked wall region.

This is a visualization/data bootstrap, not a lighting-robust learned model.
"""
import argparse
import json
from pathlib import Path

import cv2
import numpy as np


def points_arg(value):
    try:
        pts = [[int(v) for v in pair.split(",")] for pair in value.split()]
        if len(pts) < 3 or any(len(p) != 2 for p in pts):
            raise ValueError
        return np.asarray(pts, np.int32)
    except ValueError as exc:
        raise argparse.ArgumentTypeError("expected at least 3 x,y points separated by spaces") from exc


def main():
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument("image", type=Path)
    ap.add_argument("--paint-hsv", type=int, nargs=6, required=True,
                    metavar=("H_MIN", "S_MIN", "V_MIN", "H_MAX", "S_MAX", "V_MAX"))
    ap.add_argument("--wall", type=points_arg, required=True, help='wall polygon: "x,y x,y x,y ..."')
    ap.add_argument("--out", type=Path, default=Path("painting_overlay.png"))
    ap.add_argument("--rectify", action="store_true", help="rectify a four-corner wall to a front view")
    args = ap.parse_args()

    image = cv2.imread(str(args.image))
    if image is None:
        ap.error(f"could not read image: {args.image}")
    h, w = image.shape[:2]
    if np.any(args.wall[:, 0] < 0) or np.any(args.wall[:, 0] >= w) or np.any(args.wall[:, 1] < 0) or np.any(args.wall[:, 1] >= h):
        ap.error("wall polygon points must fall inside the image")
    if args.rectify:
        if len(args.wall) != 4:
            ap.error("--rectify requires exactly four wall corners, ordered clockwise from top-left")
        # Rectified dimensions preserve approximate source edge lengths.
        p = args.wall.astype(np.float32)
        width = max(int(np.linalg.norm(p[1] - p[0])), int(np.linalg.norm(p[2] - p[3])))
        height = max(int(np.linalg.norm(p[3] - p[0])), int(np.linalg.norm(p[2] - p[1])))
        dst = np.float32([[0, 0], [width - 1, 0], [width - 1, height - 1], [0, height - 1]])
        image = cv2.warpPerspective(image, cv2.getPerspectiveTransform(p, dst), (width, height))
        wall_mask = np.full((height, width), 255, np.uint8)
    else:
        wall_mask = np.zeros((h, w), np.uint8)
        cv2.fillPoly(wall_mask, [args.wall], 255)

    lo = np.array(args.paint_hsv[:3], dtype=np.uint8)
    hi = np.array(args.paint_hsv[3:], dtype=np.uint8)
    hsv = cv2.cvtColor(image, cv2.COLOR_BGR2HSV)
    paint = cv2.inRange(hsv, lo, hi)
    paint = cv2.bitwise_and(paint, wall_mask)
    wall_pixels = int(cv2.countNonZero(wall_mask))
    painted_pixels = int(cv2.countNonZero(paint))
    coverage = painted_pixels / wall_pixels if wall_pixels else 0.0

    overlay = image.copy()
    overlay[paint > 0] = (0, 220, 0)
    overlay[wall_mask == 0] = (overlay[wall_mask == 0] * 0.35).astype(np.uint8)
    result = cv2.addWeighted(image, 0.58, overlay, 0.42, 0)
    cv2.polylines(result, [args.wall] if not args.rectify else [np.array([[0, 0], [w-1, 0], [w-1, h-1], [0, h-1]])], True, (0, 255, 255), 2)
    args.out.parent.mkdir(parents=True, exist_ok=True)
    if not cv2.imwrite(str(args.out), result):
        ap.error(f"could not write output: {args.out}")
    print(json.dumps({"painted_fraction": coverage, "painted_percent": round(coverage * 100, 2),
                      "painted_pixels": painted_pixels, "wall_pixels": wall_pixels,
                      "uncertainty_fraction": None,
                      "note": "HSV baseline; unclassified pixels are not yet distinguished from unpainted pixels."}, indent=2))
    print(f"overlay: {args.out}")


if __name__ == "__main__":
    main()