hongchen/Floorplan/full_server.py

"""
Full Floorplan Pipeline - Single Endpoint Server.
Call once, everything is handled: RGB → initial_mask → refine_mask → JSON

Endpoints:
  POST /process  — run the complete pipeline for a folder
  GET  /health   — health check
"""

import os
import gc
import json
import subprocess
import sys

import cv2
import numpy as np
import shutil
from fastapi import FastAPI, Form, UploadFile, File
from pydantic import BaseModel, Field

app = FastAPI(title="Floorplan Pipeline")

# ========================================================================
# 1. extract_initial_mask (BiRefNet)
# ========================================================================

def red_edge_generate(img_path, save_path):
    img = cv2.imread(img_path)
    img_2 = np.zeros_like(img)
    mask = (img[:, :, 0] == 0) * (img[:, :, 1] == 0) * (img[:, :, 2] == 0)
    img_2[~mask] = (255, 255, 255)
    edges = cv2.Canny(img_2, 50, 150)
    kernel = np.ones((3, 3), np.uint8)
    edges = cv2.dilate(edges, kernel, 1)
    mask = edges[:, :, None] / 255.0
    masks = np.concatenate([mask, mask, mask], axis=-1)
    img1 = (masks * (0.0, 0.0, 255.0)).clip(0, 255)
    alpha = 0.9
    img = img1 * alpha + img * (1 - masks * alpha)
    cv2.imwrite(save_path, img)


def generate_initial_mask(rgb_path, output_path, model_path, img_name):
    """Subprocess: red-edge + BiRefNet → initial_mask"""
    result = subprocess.run([
        sys.executable, "-c", f'''
import os, gc, torch, onnxruntime, cv2, numpy as np
from PIL import Image

def red_edge_generate(img_path, save_path):
    img = cv2.imread(img_path)
    img_2 = np.zeros_like(img)
    mask = (img[:, :, 0] == 0) * (img[:, :, 1] == 0) * (img[:, :, 2] == 0)
    img_2[~mask] = (255, 255, 255)
    edges = cv2.Canny(img_2, 50, 150)
    kernel = np.ones((3, 3), np.uint8)
    edges = cv2.dilate(edges, kernel, 1)
    mask = edges[:, :, None] / 255.0
    masks = np.concatenate([mask, mask, mask], axis=-1)
    img1 = (masks * (0.0, 0.0, 255.0)).clip(0, 255)
    alpha = 0.9
    img = img1 * alpha + img * (1 - masks * alpha)
    cv2.imwrite(save_path, img)

def predict_birefnet_onnx(image_path, onnx_session, mask_dir, input_size=(1024, 1024)):
    orig_img = Image.open(image_path).convert("RGB")
    w_orig, h_orig = orig_img.size
    img_resized = orig_img.resize(input_size, resample=Image.BILINEAR)
    img_np = np.array(img_resized).astype(np.float32) / 255.0
    mean = np.array([0.485, 0.456, 0.406], dtype=np.float32)
    std = np.array([0.229, 0.224, 0.225], dtype=np.float32)
    img_np = (img_np - mean) / std
    img_np = img_np.transpose(2, 0, 1)[np.newaxis, :]
    img_np = np.ascontiguousarray(img_np)
    input_name = onnx_session.get_inputs()[0].name
    outputs = onnx_session.run(None, {{input_name: img_np}})
    raw_preds = outputs[-1]
    pred_mask = 1 / (1 + np.exp(-raw_preds))
    pred_mask = pred_mask.squeeze()
    mask_resized = cv2.resize(pred_mask, (w_orig, h_orig), interpolation=cv2.INTER_LINEAR)
    mask_8bit = (mask_resized * 255).astype(np.uint8)
    kernel = np.ones((3, 3), np.uint8)
    mask_eroded = cv2.erode(mask_8bit, kernel, iterations=1)
    if not os.path.exists(mask_dir):
        os.makedirs(mask_dir)
    save_path = os.path.join(mask_dir, os.path.basename(image_path))
    cv2.imwrite(save_path, mask_eroded)

red_edge_generate("{rgb_path}", "{rgb_path}")
session = onnxruntime.InferenceSession("{model_path}", providers=[("CUDAExecutionProvider", {{"device_id": 0}})])
temp_dir = "{os.path.dirname(output_path)}/_temp_masks"
os.makedirs(temp_dir, exist_ok=True)
predict_birefnet_onnx("{rgb_path}", session, temp_dir)
temp_path = os.path.join(temp_dir, os.path.basename("{img_name}"))
import shutil
shutil.move(temp_path, "{output_path}")
try:
    os.rmdir(temp_dir)
except:
    pass
del session
torch.cuda.empty_cache()
torch.cuda.synchronize()
gc.collect()
gc.collect()
''',
    ], capture_output=True, text=True)
    if result.returncode != 0:
        raise RuntimeError(f"initial_mask 生成失败: {result.stderr}")


# ========================================================================
# 2. inference_refine_mask (FLUX.2-klein)
# ========================================================================

def generate_refine_mask(initial_mask_path, output_path, flux_model_path):
    """Subprocess: FLUX.2-klein initial_mask → refine_mask"""
    result = subprocess.run([
        sys.executable, "-c", f'''
import gc, torch
import os as _os
_os.environ["CUDA_VISIBLE_DEVICES"] = "1"
from diffusers import Flux2KleinPipeline
from diffusers.utils import load_image

pipe = Flux2KleinPipeline.from_pretrained("{flux_model_path}", torch_dtype=torch.bfloat16)
pipe = pipe.to("cuda")
generator = torch.Generator(device="cuda").manual_seed(0)

image = load_image("{initial_mask_path}")
base_width, base_height = image.size
target_width = (base_width // 8) * 8
target_height = (base_height // 8) * 8

prompt = """
(best quality, 4k), architectural floor plan mask, instance segmentation,
do not add extra blocks,
distinct separate white blocks, clear black gaps between rooms,
separated connected components, clean sharp edges, top-down view,
binary mask style, white rooms on black background, no touching blocks,
The image should be positioned exactly as it was in the original image; do not shift it.
logical room separation
"""

result_img = pipe(
    image=image, prompt=prompt,
    height=target_height, width=target_width,
    guidance_scale=4.0, num_inference_steps=4,
    generator=generator
).images[0]

result_img.resize((base_width, base_height)).save("{output_path}")

del pipe, generator
torch.cuda.empty_cache()
torch.cuda.synchronize()
gc.collect()
gc.collect()
''',
    ], capture_output=True, text=True, timeout=300)
    if result.returncode != 0:
        raise RuntimeError(f"refine_mask 生成失败: {result.stderr}")


# ========================================================================
# 3. pipeline (core processing)
# ========================================================================

_model_cache = {}

def _load_yolo(p):
    if p not in _model_cache:
        from ultralytics import YOLO
        _model_cache[p] = YOLO(p)
    return _model_cache[p]


def _to_gray(mask):
    return mask if len(mask.shape) == 2 else cv2.cvtColor(mask, cv2.COLOR_BGR2GRAY)


def remove_pure_red(img):
    if img is None:
        return
    red = (img[:, :, 0] == 0) & (img[:, :, 1] == 0) & (img[:, :, 2] == 255)
    img[red] = [0, 0, 0]
    return img


def remove_edge_regions_image(img):
    result = img.copy()
    img_2 = np.zeros_like(result)
    mask = (result[:, :, 0] == 0) & (result[:, :, 1] == 0) & (result[:, :, 2] == 0)
    img_2[~mask] = (255, 255, 255)
    edges = cv2.Canny(img_2, 50, 150)
    kernel = np.ones((9, 9), np.uint8)
    edges = cv2.dilate(edges, kernel, 1)
    result[edges > 0] = (0, 0, 0)
    return remove_pure_red(result)


def extract_gaps_from_mask(mask):
    _, binary = cv2.threshold(mask, 127, 255, cv2.THRESH_BINARY)
    kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (25, 25))
    stitched = cv2.morphologyEx(binary, cv2.MORPH_CLOSE, kernel)
    gaps = cv2.subtract(stitched, binary)
    rk = cv2.getStructuringElement(cv2.MORPH_RECT, (3, 3))
    gaps_d = cv2.dilate(gaps, rk, 1)
    return gaps, cv2.add(mask, gaps_d)


def extract_mask_region_from_arrays(rgb, ori, full):
    _, bo = cv2.threshold(ori, 127, 255, cv2.THRESH_BINARY)
    _, bf = cv2.threshold(full, 127, 255, cv2.THRESH_BINARY)
    r1 = cv2.bitwise_and(rgb, rgb, mask=bo)
    r2 = cv2.bitwise_and(rgb, rgb, mask=bf)
    return cv2.subtract(r2, r1)


def _expand_rect(x, y, w, h, ep, wmax, hmax):
    if ep <= 0:
        return int(x), int(y), int(w), int(h)
    return (max(0, int(x) - ep), max(0, int(y) - ep),
            max(1, min(wmax, int(x) + int(w) + ep) - max(0, int(x) - ep)),
            max(1, min(hmax, int(y) + int(h) + ep) - max(0, int(y) - ep)))


def merge_gap_fillers_from_arrays(m1, m2, image_path="", dilation_kernel_size=5,
                                   center_threshold=50, expand_pixel=10, min_rect_short_side=30):
    m1 = _to_gray(m1)
    m2 = _to_gray(m2)
    if m1.shape != m2.shape:
        m2 = cv2.resize(m2, (m1.shape[1], m1.shape[0]))
    _, m1b = cv2.threshold(m1, 127, 255, cv2.THRESH_BINARY)
    _, m2b = cv2.threshold(m2, 0, 255, cv2.THRESH_BINARY)
    num1, bl = cv2.connectedComponents(m1b)
    num2, fl, _, _ = cv2.connectedComponentsWithStats(m2b)
    result = cv2.cvtColor(m1b, cv2.COLOR_GRAY2BGR)
    bridge = []
    for i in range(1, num2):
        sfm = (fl == i).astype(np.uint8) * 255
        k = np.ones((dilation_kernel_size, dilation_kernel_size), np.uint8)
        df = cv2.dilate(sfm, k, 1)
        tl = np.unique(bl[df > 0])
        nb = sorted(int(n) for n in tl if n > 0)
        if len(nb) == 2:
            pts = np.column_stack(np.where(sfm > 0))
            if len(pts) > 0:
                cy, cx = np.mean(pts, axis=0)
                bridge.append({'id': i, 'cx': cx, 'cy': cy,
                               'bp': tuple(int(n - 1) for n in nb), 'mask': sfm})
    groups = []
    for frag in bridge:
        assigned = False
        for g in groups:
            if g[0]['bp'] != frag['bp']:
                continue
            for ex in g:
                if abs(frag['cx'] - ex['cx']) < center_threshold or abs(frag['cy'] - ex['cy']) < center_threshold:
                    g.append(frag)
                    assigned = True
                    break
            if assigned:
                break
        if not assigned:
            groups.append([frag])
    areas, rid = [], 0
    h, w = m1.shape[:2]
    for group in groups:
        plist, frects = [], []
        for frag in group:
            pts = np.column_stack(np.where(frag['mask'] > 0))
            if pts.size > 0:
                px = pts[:, ::-1]
                plist.append(px)
                rx, ry, rw, rh = cv2.boundingRect(px)
                frects.append((int(rx), int(ry), int(rw), int(rh), frag))
        if len(plist) < 1:
            continue
        ap = np.vstack(plist)
        if len(ap) < 3:
            continue
        rx, ry, rw, rh = cv2.boundingRect(ap)
        mss = min(rw, rh)
        if len(group) > 1 and mss > min_rect_short_side:
            for fx, fy, fw, fh, frag in frects:
                if min(fw, fh) <= min_rect_short_side:
                    fx, fy, fw, fh = _expand_rect(fx, fy, fw, fh, expand_pixel, w, h)
                    cv2.rectangle(result, (fx, fy), (fx + fw, fy + fh), (0, 255, 0), -1)
                    if min(fw, fh) >= 25:
                        areas.append({'id': rid, 'x': fx, 'y': fy, 'w': fw, 'h': fh,
                                      'block_pair': [int(n) for n in frag['bp']], 'label': 'door'})
                        rid += 1
            continue
        if mss <= min_rect_short_side:
            x, y, ww, hh = _expand_rect(rx, ry, rw, rh, expand_pixel, w, h)
            cv2.rectangle(result, (x, y), (x + ww, y + hh), (0, 255, 0), -1)
            if min(ww, hh) >= 25:
                areas.append({'id': rid, 'x': x, 'y': y, 'w': ww, 'h': hh,
                              'block_pair': [int(n) for n in group[0]['bp']], 'label': 'door'})
                rid += 1
    return {
        'image_path': str(image_path),
        'image_size': {'width': int(m1.shape[1]), 'height': int(m1.shape[0])},
        'connect_area': areas
    }, result, {
        'mask1_blocks': num1 - 1, 'mask2_fragments': num2 - 1,
        'bridge_fragments': len(bridge), 'group_count': len(groups),
        'connect_area_count': len(areas)
    }


def build_block_data(rgb_img, block_mask, model_path="room_cls.pt"):
    model = _load_yolo(model_path)
    blocks = _to_gray(block_mask)
    if rgb_img.shape[:2] != blocks.shape[:2]:
        rgb_img = cv2.resize(rgb_img, (blocks.shape[1], blocks.shape[0]))
    _, bb = cv2.threshold(blocks, 127, 255, cv2.THRESH_BINARY)
    num_blocks, bl = cv2.connectedComponents(bb, connectivity=8)
    blist = []
    for b in range(1, num_blocks):
        ms = (bl == b).astype(np.uint8)
        cnt, _ = cv2.findContours(ms, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        pts, cx, cy = [], 0, 0
        if cnt:
            lc = max(cnt, key=cv2.contourArea)
            simp = cv2.approxPolyDP(lc, 2.0, True)
            for pt in simp:
                pts.extend([int(pt[0][0]), int(pt[0][1])])
            M = cv2.moments(simp)
            if M['m00']:
                cx, cy = int(M['m10'] / M['m00']), int(M['m01'] / M['m00'])
            else:
                cx, cy = int(np.mean(simp[:, 0, 0])), int(np.mean(simp[:, 0, 1]))
        x, y, w, h = cv2.boundingRect(ms)
        if not cnt:
            cx, cy = int(x + w / 2), int(y + h / 2)
        pad = 20
        y1, y2 = max(0, y - pad), min(rgb_img.shape[0], y + h + pad)
        x1, x2 = max(0, x - pad), min(rgb_img.shape[1], x + w + pad)
        roi = rgb_img[y1:y2, x1:x2].copy()
        roi[ms[y1:y2, x1:x2] == 0] = [0, 0, 0]
        label, conf = "other_room", 0.0
        if roi.shape[0] > 10 and roi.shape[1] > 10:
            res = model(roi, verbose=False)[0]
            if hasattr(res, 'probs') and res.probs is not None:
                tc = float(res.probs.top1conf.cpu().numpy())
                if tc >= 0.15:
                    label, conf = model.names[int(res.probs.top1)], tc
        blist.append({"id": b - 1, "points": pts, "label": label, "center": [cx, cy]})
    return blist


def detect_furniture(rgb_img, model_path='furniture_detect.onnx'):
    model = _load_yolo(model_path)
    res = model(rgb_img, conf=0.25, verbose=False)[0]
    allowed = {'sofa', 'chair', 'desk', 'bed', 'window'}
    fl = []
    if len(res.boxes) > 0:
        bx = res.boxes.xyxy.cpu().numpy()
        bc = res.boxes.cls.cpu().numpy()
        for i in range(len(bx)):
            b1, y1, b2, y2 = [int(v) for v in bx[i]]
            lb = model.names[int(bc[i])]
            if lb not in allowed:
                continue
            fl.append({'id': len(fl), 'label': lb,
                       'center': [(b1 + b2) // 2, (y1 + y2) // 2],
                       'points': {'x1': b1, 'y1': y1, 'x2': b2, 'y2': y1,
                                  'x3': b2, 'y3': y2, 'x4': b1, 'y4': y2}})
    return fl


# ── Refinement / normalization / merge (inline from pipeline.py) ────────

def refine_blocks_in_data(data):
    blocks = data.get("block", [])
    total = 0
    for idx, block in enumerate(blocks):
        pd2 = block.get("points", [])
        if not pd2:
            continue
        if isinstance(pd2[0], list) and len(pd2[0]) == 4:
            segs = [[float(v) for v in s] for s in pd2]
        else:
            pts = [[float(pd2[i]), float(pd2[i + 1])] for i in range(0, len(pd2), 2) if i + 1 < len(pd2)]
            segs = [[p[0], p[1], pts[(j + 1) % len(pts)][0], pts[(j + 1) % len(pts)][1]] for j, p in enumerate(pts)]
        refined = refine_single_block_segments(segs) or segs
        segs_int = [[int(round(s[0])), int(round(s[1])), int(round(s[2])), int(round(s[3]))] for s in refined]
        block["points"] = segs_int
        block["refined"] = True
        block["format"] = "segments"
        block["segment_count"] = len(segs_int)
        total += len(segs_int)
    data["format_version"] = "segments_v1"
    data["total_segments"] = total


def refine_single_block_segments(segments):
    if not segments:
        return []
    g = orthogonalize_and_move_nodes([segments], 15)[0]
    r = apply_user_refinement(g, 30, 30)
    r = merge_collinear(r, 2)
    r = merge_parallel(r, 12)
    return r


def to_key(p):
    return (round(float(p[0]), 1), round(float(p[1]), 1))


def is_orthogonal(seg, t=1e-1):
    return abs(seg[2] - seg[0]) < t or abs(seg[3] - seg[1]) < t


def orthogonalize_and_move_nodes(sgs, ath):
    out = []
    for grp in sgs:
        np2 = {}
        def gsn(p):
            pk = to_key(p)
            for ep, o in np2.items():
                if np.linalg.norm(np.array(pk) - np.array(ep)) < 2.5:
                    return o
            nn = np.array(p, dtype=np.float32)
            np2[pk] = nn
            return nn
        gs = [(gsn(s[:2]), gsn(s[2:])) for s in grp]
        for _ in range(3):
            for p1, p2 in gs:
                dx, dy = abs(p2[0] - p1[0]), abs(p2[1] - p1[1])
                a = np.degrees(np.arctan2(dy, dx))
                if a < ath or a > (180 - ath):
                    ay = (p1[1] + p2[1]) / 2
                    p1[1] = p2[1] = ay
                elif abs(a - 90) < ath:
                    ax = (p1[0] + p2[0]) / 2
                    p1[0] = p2[0] = ax
        out.append([[p1[0], p1[1], p2[0], p2[1]] for p1, p2 in gs])
    return out


def apply_user_refinement(grp, lt=30, at=30):
    if len(grp) < 2:
        return grp
    atr = np.radians(at)
    si = next((i for i, s in enumerate(grp) if is_orthogonal(s)), -1)
    if si == -1:
        return grp
    wl = [list(s) for s in (grp[si:] + grp[:si])]
    ref, i = [], 0
    while i < len(wl):
        cs = wl[i]
        ref.append(cs)
        ni = i + 1
        if ni >= len(wl):
            break
        if not is_orthogonal(wl[ni]):
            flp, tgl, gs = False, 0, []
            pi = ni
            while pi < len(wl):
                if is_orthogonal(wl[pi]):
                    flp = True
                    break
                s = wl[pi]
                tgl += np.sqrt((s[2] - s[0]) ** 2 + (s[3] - s[1]) ** 2)
                gs.append(s)
                pi += 1
            if flp:
                lp = wl[pi]
                if tgl < lt:
                    pe, ls = [cs[2], cs[3]], [lp[0], lp[1]]
                    ih1, ih2 = abs(cs[3] - cs[1]) < 1e-1, abs(lp[3] - lp[1]) < 1e-1
                    pm = [pe[0], ls[1]] if ih1 == ih2 else [ls[0], pe[1]] if ih1 else [pe[0], ls[1]]
                    lp[0], lp[1] = pm[0], pm[1]
                    ref.append([pe[0], pe[1], lp[0], lp[1]])
                    i = pi
                else:
                    gs[0][0], gs[0][1] = cs[2], cs[3]
                    mg, ts = [], list(gs[0])
                    for k in range(1, len(gs)):
                        ns = gs[k]
                        v1, v2 = (ts[2] - ts[0], ts[3] - ts[1]), (ns[2] - ns[0], ns[3] - ns[1])
                        m1, m2 = np.sqrt(v1[0]**2 + v1[1]**2), np.sqrt(v2[0]**2 + v2[1]**2)
                        if m1 > 1e-6 and m2 > 1e-6:
                            ct = abs(v1[0]*v2[0] + v1[1]*v2[1]) / (m1 * m2)
                            if np.arccos(max(-1, min(1, ct))) < atr:
                                ts[2], ts[3] = ns[2], ns[3]
                                continue
                        mg.append(ts)
                        ts = list(ns)
                        ts[0], ts[1] = mg[-1][2], mg[-1][3]
                    mg.append(ts)
                    mg[-1][2], mg[-1][3] = lp[0], lp[1]
                    ref.extend(mg)
                    i = pi
            else:
                ref.extend(wl[ni:])
                break
        else:
            i += 1
    return ref


def merge_collinear(og, dt=0.5):
    if len(og) < 2:
        return og
    mg, cs = [], list(og[0])
    for i in range(1, len(og)):
        ns = og[i]
        ih = abs(cs[1] - cs[3]) < dt and abs(ns[1] - ns[3]) < dt and abs(cs[3] - ns[1]) < dt
        iv = abs(cs[0] - cs[2]) < dt and abs(ns[0] - ns[2]) < dt and abs(cs[2] - ns[0]) < dt
        if ih or iv:
            cs[2], cs[3] = ns[2], ns[3]
        else:
            mg.append(cs)
            cs = list(ns)
    mg.append(cs)
    return mg


def merge_parallel(grp, dt=15.0):
    if not grp:
        return grp
    sg = [list(s) for s in grp]
    it = 0
    while True:
        mir, mi = False, set()
        for i in range(len(sg)):
            if i in mi:
                continue
            for j in range(i + 1, len(sg)):
                if j in mi:
                    continue
                s1, s2 = sg[i], sg[j]
                ih1, ih2 = abs(s1[1] - s1[3]) < 1e-1, abs(s2[1] - s2[3]) < 1e-1
                iv1, iv2 = abs(s1[0] - s1[2]) < 1e-1, abs(s2[0] - s2[2]) < 1e-1
                if ih1 and ih2:
                    d = abs(s1[1] - s2[1])
                    ol = min(max(s1[0], s1[2]), max(s2[0], s2[2])) - max(min(s1[0], s1[2]), min(s2[0], s2[2]))
                    if d < dt and ol > 0:
                        np2 = (s1[1] + s2[1]) / 2
                        for s in sg:
                            if abs(s[1] - s1[1]) < 1e-1 or abs(s[1] - s2[1]) < 1e-1:
                                s[1] = np2
                            if abs(s[3] - s1[1]) < 1e-1 or abs(s[3] - s2[1]) < 1e-1:
                                s[3] = np2
                        s1[0], s1[2] = min(s1[0], s1[2], s2[0], s2[2]), max(s1[0], s1[2], s2[0], s2[2])
                        s1[1], s1[3] = min(s1[1], s1[3], s2[1], s2[3]), max(s1[1], s1[3], s2[1], s2[3])
                        mi.add(j)
                        mir = True
                        break
                elif iv1 and iv2:
                    d = abs(s1[0] - s2[0])
                    ol = min(max(s1[1], s1[3]), max(s2[1], s2[3])) - max(min(s1[1], s1[3]), min(s2[1], s2[3]))
                    if d < dt and ol > 0:
                        np2 = (s1[0] + s2[0]) / 2
                        for s in sg:
                            if abs(s[0] - s1[0]) < 1e-1 or abs(s[0] - s2[0]) < 1e-1:
                                s[0] = np2
                            if abs(s[2] - s1[0]) < 1e-1 or abs(s[2] - s2[0]) < 1e-1:
                                s[2] = np2
                        s1[0], s1[2] = min(s1[0], s1[2], s2[0], s2[2]), max(s1[0], s1[2], s2[0], s2[2])
                        s1[1], s1[3] = min(s1[1], s1[3], s2[1], s2[3]), max(s1[1], s1[3], s2[1], s2[3])
                        mi.add(j)
                        mir = True
                        break
            if mir:
                break
        if mi:
            sg = [s for idx, s in enumerate(sg) if idx not in mi]
        if not mir:
            break
        it += 1
        if it > 100:
            break
    return sg


def normalize_segment(seg, angle=10):
    x1, y1, x2, y2 = seg
    dx, dy = x2 - x1, y2 - y1
    if dx == 0 and dy == 0:
        return seg
    theta = abs(np.degrees(np.arctan2(abs(dy), abs(dx))))
    if theta <= angle:
        my = round((y1 + y2) / 2)
        return [x1, my, x2, my]
    if theta >= 90 - angle:
        mx = round((x1 + x2) / 2)
        return [mx, y1, mx, y2]
    return seg


def normalize_all(data, angle=10):
    for b in data.get('block', []):
        b['points'] = [normalize_segment(s, angle) for s in b['points']]


def segment_orientation(seg):
    return 'H' if seg[1] == seg[3] else ('V' if seg[0] == seg[2] else None)


def segments_overlap(sa, sb, ori):
    if ori == 'H':
        a1, a2, b1, b2 = min(sa[0], sa[2]), max(sa[0], sa[2]), min(sb[0], sb[2]), max(sb[0], sb[2])
    else:
        a1, a2, b1, b2 = min(sa[1], sa[3]), max(sa[1], sa[3]), min(sb[1], sb[3]), max(sb[1], sb[3])
    return a1 <= b2 and b1 <= a2


def cluster_and_merge(segs, ori, thresh):
    idxs = [i for i, s in enumerate(segs) if segment_orientation(s) == ori]
    if not idxs:
        return False
    parent = {i: i for i in idxs}
    def find(x):
        while parent[x] != x:
            parent[x] = parent[parent[x]]
            x = parent[x]
        return x
    def union(x, y):
        parent[find(x)] = find(y)
    for ii in range(len(idxs)):
        for jj in range(ii + 1, len(idxs)):
            i, j = idxs[ii], idxs[jj]
            si, sj = segs[i], segs[j]
            ci = si[1] if ori == 'H' else si[0]
            cj = sj[1] if ori == 'H' else sj[0]
            if abs(ci - cj) <= thresh and segments_overlap(si, sj, ori):
                union(i, j)
    from collections import defaultdict
    clusters = defaultdict(list)
    for i in idxs:
        clusters[find(i)].append(i)
    changed = False
    for members in clusters.values():
        if len(members) < 2:
            continue
        coords = [segs[i][1] if ori == 'H' else segs[i][0] for i in members]
        ov = list(set(coords))
        if len(ov) == 1:
            continue
        nv = round(sum(coords) / len(coords))
        os2 = set(ov)
        for k, s in enumerate(segs):
            so = segment_orientation(s)
            if ori == 'H':
                if so == 'H' and s[1] in os2:
                    segs[k][1] = segs[k][3] = nv
                elif so != 'H':
                    if s[1] in os2:
                        segs[k][1] = nv
                    if s[3] in os2:
                        segs[k][3] = nv
            else:
                if so == 'V' and s[0] in os2:
                    segs[k][0] = segs[k][2] = nv
                elif so != 'V':
                    if s[0] in os2:
                        segs[k][0] = nv
                    if s[2] in os2:
                        segs[k][2] = nv
        changed = True
    return changed


def merge_all_blocks(data, threshold=6):
    blocks = data.get('block', [])
    all_segs, counts = [], []
    for b in blocks:
        ss = [list(s) for s in b['points']]
        all_segs.extend(ss)
        counts.append(len(ss))
    ch = True
    while ch:
        ch = cluster_and_merge(all_segs, 'H', threshold) or cluster_and_merge(all_segs, 'V', threshold)
    idx = 0
    seen = set()
    for b, cnt in zip(blocks, counts):
        bs = []
        for s in all_segs[idx:idx + cnt]:
            k = (min((s[0], s[1]), (s[2], s[3])), max((s[0], s[1]), (s[2], s[3])))
            if not (s[0] == s[2] and s[1] == s[3]) and k not in seen:
                seen.add(k)
                bs.append(s)
        b['points'] = bs
        b['segment_count'] = len(bs)
        idx += cnt
    data['total_segments'] = sum(len(x['points']) for x in blocks)


# ── Core pipeline ────────────────────────────────────────────────────────

def run_pipeline(rgb_img, block_mask, room_model, furniture_model,
                 merge_threshold, angle, dilation_kernel_size,
                 center_threshold, expand_pixel, min_rect_short_side):
    clean = remove_edge_regions_image(rgb_img)
    _, gap_mask = extract_gaps_from_mask(block_mask)
    conn = extract_mask_region_from_arrays(clean, block_mask, gap_mask)
    jd, _, stats = merge_gap_fillers_from_arrays(
        block_mask, conn, "", dilation_kernel_size,
        center_threshold, expand_pixel, min_rect_short_side)
    jd["connect_area_stats"] = stats
    jd["block"] = build_block_data(rgb_img, block_mask, room_model)
    jd["furniture"] = detect_furniture(rgb_img, furniture_model)
    refine_blocks_in_data(jd)
    normalize_all(jd, angle)
    merge_all_blocks(jd, merge_threshold)
    return jd


# ========================================================================
# Endpoints
# ========================================================================

class ProcessRequest(BaseModel):
    folder: str = Field(..., description="包含 RGB 图片的文件夹路径")
    rgb_pattern: str | None = Field(None, description="RGB 文件名前缀（默认用文件夹名）")
    initial_model: str = Field("initial_mask.onnx")
    refine_model: str = Field("./black-forest-labs/FLUX.2-klein-4B")
    room_model: str = Field("room_cls.pt")
    furniture_model: str = Field("furniture_detect.onnx")
    merge_threshold: int = Field(15, ge=1)
    angle: int = Field(10, ge=0, le=45)
    dilation_kernel_size: int = Field(5, ge=1)
    center_threshold: int = Field(15, ge=1)
    expand_pixel: int = Field(2, ge=0)
    min_rect_short_side: int = Field(30, ge=1)


@app.get("/health")
def health():
    return {"status": "ok"}


@app.post("/process")
def process(req: ProcessRequest):
    """完整管线：RGB → initial_mask → refine_mask → JSON"""
    folder = req.folder
    if not os.path.isdir(folder):
        return {"success": False, "error": f"文件夹不存在: {folder}"}

    # Find RGB
    prefix = req.rgb_pattern or os.path.basename(folder.rstrip("/"))
    rgbs = [f for f in os.listdir(folder)
            if f.startswith(prefix)
            and not f.startswith("initial_mask_")
            and not f.startswith("refine_mask_")
            and f.lower().endswith((".png", ".jpg", ".jpeg"))]
    if not rgbs:
        return {"success": False, "error": f"未找到 RGB 图片 (前缀={prefix})"}

    rgb_name = rgbs[0]
    rgb_path = os.path.join(folder, rgb_name)
    base_name = os.path.splitext(rgb_name)[0]

    # All outputs go into a subfolder named after the image
    out_dir = os.path.join(folder, base_name)
    os.makedirs(out_dir, exist_ok=True)

    # Copy RGB into out_dir if not already there
    rgb_in_out = os.path.join(out_dir, rgb_name)
    if not os.path.exists(rgb_in_out):
        shutil.copy2(rgb_path, rgb_in_out)

    print(f"\n===== 开始处理 =====")
    print(f"RGB: {rgb_name}")
    print(f"输出目录: {out_dir}")

    try:
        # ── Stage 1: initial_mask ──────────────────────────────────
        im_path = os.path.join(out_dir, f"initial_mask_{rgb_name}")
        if os.path.exists(im_path):
            print(f"[1/3] initial_mask 已存在，跳过")
        else:
            print(f"[1/3] 生成 initial_mask ...")
            generate_initial_mask(rgb_in_out, im_path, req.initial_model, rgb_name)
            print(f"  ✓ {im_path}")

        # ── Stage 2: refine_mask ───────────────────────────────────
        rm_path = os.path.join(out_dir, f"refine_mask_{rgb_name}")
        if os.path.exists(rm_path):
            print(f"[2/3] refine_mask 已存在，跳过")
        else:
            print(f"[2/3] 生成 refine_mask ...")
            generate_refine_mask(im_path, rm_path, req.refine_model)
            print(f"  ✓ {rm_path}")

        # ── Stage 3: pipeline ──────────────────────────────────────
        print(f"[3/3] 运行管线 ...")
        rgb_img = cv2.imread(rgb_in_out)
        block_mask = cv2.imread(rm_path, cv2.IMREAD_GRAYSCALE)
        if rgb_img is None:
            return {"success": False, "error": f"无法读取 RGB: {rgb_in_out}"}
        if block_mask is None:
            return {"success": False, "error": f"无法读取掩码: {rm_path}"}

        jd = run_pipeline(
            rgb_img, block_mask,
            req.room_model, req.furniture_model,
            req.merge_threshold, req.angle,
            req.dilation_kernel_size, req.center_threshold,
            req.expand_pixel, req.min_rect_short_side,
        )

        out_json = os.path.join(out_dir, base_name + ".json")
        with open(out_json, "w", encoding="utf-8") as f:
            json.dump(jd, f, indent=2, ensure_ascii=False)

        print(f"\n===== 完成 =====")
        return {
            "success": True,
            "stats": {
                "connect_area_count": len(jd.get("connect_area", [])),
                "block_count": len(jd.get("block", [])),
                "furniture_count": len(jd.get("furniture", [])),
                "total_segments": jd.get("total_segments", 0),
                "output_folder": out_dir,
                "json_path": out_json,
                "initial_mask_path": im_path,
                "refine_mask_path": rm_path,
            },
        }

    except Exception as e:
        import traceback
        traceback.print_exc()
        return {"success": False, "error": str(e)}


if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8070)