hongchen/Floorplan/full_server_v2.py

"""
Full Floorplan Pipeline - 5-Stage Server

Endpoints:
  POST /process       — local folder mode (folder + rgb_pattern)
  POST /process-upload — upload image mode (multipart/form-data)
  GET  /health        — health check

Pipeline (5 stages):
  1. extract_initial_mask (BiRefNet)
  2. inference_refine_mask (FLUX.2-klein)
  3. pipeline → JSON
  4. vis → visualization PNG
  5. pixel_to_world → demo.json (world coordinates)
"""

import os
import json
import subprocess
import sys
import uuid
import tempfile

import cv2
import numpy as np
import shutil
from pixel_to_world import (
    build_vertices_and_segments as _build_vertices_and_segments,
    build_shapes as _build_shapes_world,
)
from fastapi import FastAPI, UploadFile, File
from fastapi.responses import JSONResponse
from pydantic import BaseModel, Field

app = FastAPI(title="Floorplan Pipeline")


# ========================================================================
# Stage 1: extract_initial_mask (BiRefNet)
# ========================================================================

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, timeout=300)
    if result.returncode != 0:
        raise RuntimeError(f"initial_mask 生成失败: {result.stderr}")


# ========================================================================
# Stage 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
"""

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

pipe_result.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}")


# ========================================================================
# Stage 3: pipeline (core processing) - inline from pipeline.py
# ========================================================================

_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 ──────────────────────────────────

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
    print(f"[DEBUG] Stage3门检测: bridge_fragments={stats.get('bridge_fragments', 0)}, "
          f"group_count={stats.get('group_count', 0)}, "
          f"connect_area_count={stats.get('connect_area_count', 0)}, "
          f"mask1_blocks={stats.get('mask1_blocks', 0)}, "
          f"mask2_fragments={stats.get('mask2_fragments', 0)}")
    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


# ========================================================================
# Stage 4: vis - inline from vis.py
# ========================================================================

COLOR_MAP = {
    "living_room": (255, 180, 100),
    "bed_room": (100, 255, 100),
    "bath_room": (255, 100, 255),
    "kitchen_room": (100, 255, 255),
    "other_room": (180, 180, 180),
    "balcony": (200, 150, 100),
}
DEFAULT_COLOR = (200, 200, 200)


def get_image_size(data):
    if 'image_size' in data:
        return data['image_size']['width'], data['image_size']['height']
    max_x, max_y = 0, 0
    for block in data.get('block', []):
        segments = block.get('points', [])
        if isinstance(segments, list) and len(segments) > 0:
            if isinstance(segments[0], list) and len(segments[0]) == 4:
                for seg in segments:
                    max_x = max(max_x, seg[0], seg[2])
                    max_y = max(max_y, seg[1], seg[3])
    for area in data.get('connect_area', []):
        max_x = max(max_x, area.get('x', 0) + area.get('w', 0))
        max_y = max(max_y, area.get('y', 0) + area.get('h', 0))
    return int(max_x) + 100, int(max_y) + 100


def get_points_from_segments(segments):
    if not segments:
        return []
    points_set = set()
    for seg in segments:
        if len(seg) == 4:
            points_set.add((int(seg[0]), int(seg[1])))
            points_set.add((int(seg[2]), int(seg[3])))
    return list(points_set)


def visualize_json(json_data, output_path, rgb_path=None, vis_door=False):
    """从内存 JSON 数据绘制可视化 PNG。"""
    width, height = get_image_size(json_data)

    if rgb_path and os.path.exists(rgb_path):
        canvas = cv2.imread(rgb_path)
    else:
        canvas = np.zeros((height, width, 3), dtype=np.uint8)

    blocks = json_data.get('block', [])
    for block_idx, block in enumerate(blocks):
        block_id = block.get('id', block_idx)
        label = block.get('label', 'door')
        segments = block.get('points', [])
        center = block.get('center', None)
        color = COLOR_MAP.get(label.lower(), DEFAULT_COLOR)

        if not (isinstance(segments, list) and len(segments) > 0):
            continue
        if not (isinstance(segments[0], list) and len(segments[0]) == 4):
            continue

        points = get_points_from_segments(segments)
        if len(points) < 2:
            continue

        for seg in segments:
            p1 = (int(seg[0]), int(seg[1]))
            p2 = (int(seg[2]), int(seg[3]))
            cv2.line(canvas, p1, p2, color, 2, cv2.LINE_AA)

        for pt in points:
            cv2.circle(canvas, pt, 4, (255, 255, 255), -1)

        if center and len(center) == 2:
            cx, cy = int(center[0]), int(center[1])
        else:
            pts = np.array(points, dtype=np.int32)
            M = cv2.moments(pts)
            if M["m00"] != 0:
                cx = int(M["m10"] / M["m00"])
                cy = int(M["m01"] / M["m00"])
            else:
                cx, cy = int(np.mean(pts[:, 0])), int(np.mean(pts[:, 1]))

        font = cv2.FONT_HERSHEY_SIMPLEX
        font_scale = 0.6
        thickness = 1
        (tw, th), _ = cv2.getTextSize(label, font, font_scale, thickness)
        cx = max(tw // 2 + 5, min(cx, width - tw // 2 - 5))
        cy = max(th // 2 + 5, min(cy, height - th // 2 - 5))

        cv2.rectangle(canvas, (cx - tw // 2 - 3, cy - th - 3),
                      (cx + tw // 2 + 3, cy + 3), (255, 255, 255), -1)
        cv2.putText(canvas, label, (cx - tw // 2, cy),
                    font, font_scale, (0, 0, 0), thickness)

        id_text = f"ID:{block_id}"
        (iw, ih), _ = cv2.getTextSize(id_text, font, 0.4, 1)
        cv2.putText(canvas, id_text, (cx - iw // 2, cy + 15),
                    font, 0.4, (100, 100, 100), 1)

    connect_areas = json_data.get('connect_area', [])
    for area_idx, area in enumerate(connect_areas):
        x = area.get('x', 0)
        y = area.get('y', 0)
        w = area.get('w', 0)
        h = area.get('h', 0)
        label = area.get('label', 'door').lower()
        x = max(0, min(x, width - 1))
        y = max(0, min(y, height - 1))
        w = max(1, min(w, width - x))
        h = max(1, min(h, height - y))
        if vis_door and "door" in label:
            cv2.rectangle(canvas, (x, y), (x + w, y + h), (0, 0, 255), -1)
            cv2.rectangle(canvas, (x, y), (x + w, y + h), (255, 255, 255), 1)

    furniture_list = json_data.get('furniture', [])
    FURNITURE_COLOR = (0, 165, 255)
    for item in furniture_list:
        label = item.get('label', '')
        center = item.get('center', None)
        pts = item.get('points', {})
        if pts:
            bx1, by1 = pts['x1'], pts['y1']
            bx2, by2 = pts['x3'], pts['y3']
            cv2.rectangle(canvas, (bx1, by1), (bx2, by2), FURNITURE_COLOR, 2)
        if center and len(center) == 2:
            cx, cy = int(center[0]), int(center[1])
            font = cv2.FONT_HERSHEY_SIMPLEX
            (tw, th), _ = cv2.getTextSize(label, font, 0.5, 1)
            cv2.rectangle(canvas, (cx - tw // 2 - 3, cy - th - 3),
                          (cx + tw // 2 + 3, cy + 3), FURNITURE_COLOR, -1)
            cv2.putText(canvas, label, (cx - tw // 2, cy),
                        font, 0.5, (255, 255, 255), 1)

    cv2.imwrite(output_path, canvas)
    return output_path


# ========================================================================
# Stage 5: pixel_to_world - inline from pixel_to_world.py
# ========================================================================

PIXEL_SIZE = 0.01


def _pw(px, py, W, H):
    return (round(PIXEL_SIZE * (2 * px - W), 6),
            round(PIXEL_SIZE * (H - 2 * py), 6))




def convert_to_demo_format(json_data, output_path):
    """将 pipeline JSON 转换为 demo.json格式（世界坐标 + 图结构）。"""
    W = json_data["image_size"]["width"]
    H = json_data["image_size"]["height"]

    raw_blocks = []
    pixel_segments_raw = []
    seen_px_segs = set()
    for block in json_data.get("block", []):
        raw_segs = []
        pixel_polygon = []
        seen_p = set()
        for seg in block["points"]:
            wx1, wy1 = _pw(seg[0], seg[1], W, H)
            wx2, wy2 = _pw(seg[2], seg[3], W, H)
            raw_segs.append([wx1, wy1, wx2, wy2])

            # 记录像素线段（去重）用于合并
            px_key = (round(seg[0], 1), round(seg[1], 1), round(seg[2], 1), round(seg[3], 1))
            if px_key not in seen_px_segs:
                seen_px_segs.add(px_key)
                pixel_segments_raw.append((seg[0], seg[1], seg[2], seg[3]))

            for px, py in [(seg[0], seg[1]), (seg[2], seg[3])]:
                key = (round(px, 2), round(py, 2))
                if key not in seen_p:
                    seen_p.add(key)
                    pixel_polygon.append([px, py])

        world_polygon = []
        seen_w = set()
        for seg in raw_segs:
            for pt in [(seg[0], seg[1]), (seg[2], seg[3])]:
                key = (round(pt[0], 4), round(pt[1], 4))
                if key not in seen_w:
                    seen_w.add(key)
                    world_polygon.append(list(pt))

        raw_blocks.append({
            "id": block["id"],
            "label": block.get("label", "other_room"),
            "raw_segments": raw_segs,
            "raw_polygon": pixel_polygon,
            "points": world_polygon,
        })

    furniture = []
    windows_from_furniture = []
    for furn in json_data.get("furniture", []):
        raw = furn.get("points", {})
        if isinstance(raw, dict):
            pixel_pts = [(raw["x1"], raw["y1"]), (raw["x2"], raw["y2"]),
                         (raw["x3"], raw["y3"]), (raw["x4"], raw["y4"])]
        else:
            pixel_pts = raw
        world_pts = [list(_pw(p[0], p[1], W, H)) for p in pixel_pts]
        item = {
            "id": furn.get("id", 0),
            "label": furn.get("label", "unknown"),
            "points": world_pts,
            "raw_pts": pixel_pts,
            "score": furn.get("score", 0.9),
        }
        if furn.get("label", "").lower() == "window":
            windows_from_furniture.append(item)
        else:
            furniture.append(item)

    connect_areas = []
    for area in json_data.get("connect_area", []):
        world_area = {
            "id": area.get("id", 0),
            "label": area.get("label", "door"),
            "block_pair": area.get("block_pair", []),
            "_W": W, "_H": H,
        }
        if "category" in area:
            world_area["category"] = area["category"]
        if "score" in area:
            world_area["score"] = area["score"]

        if "x" in area and "w" in area:
            # 保留像素坐标用于 bbox 贯穿检测
            world_area["_px"] = area["x"]
            world_area["_py"] = area["y"]
            world_area["_pw"] = area["w"]
            world_area["_ph"] = area["h"]

            wx, wy = _pw(area["x"], area["y"], W, H)
            wx2, wy2 = _pw(area["x"] + area["w"], area["y"] + area["h"], W, H)
            world_area["x"] = wx
            world_area["y"] = wy
            world_area["w"] = abs(wx2 - wx)
            world_area["h"] = abs(wy2 - wy)
            world_area["points"] = [
                [world_area["x"], world_area["y"]],
                [world_area["x"] + world_area["w"], world_area["y"]],
                [world_area["x"] + world_area["w"], world_area["y"] + world_area["h"]],
                [world_area["x"], world_area["y"] + world_area["h"]],
            ]
        elif "points" in area:
            world_area["points"] = [list(_pw(p[0], p[1], W, H)) for p in area["points"]]
        connect_areas.append(world_area)

    door_areas = [a for a in connect_areas if a.get("label", "").lower() == "door"]
    vertices, segments, split_pixel_segments = _build_vertices_and_segments(raw_blocks, W, H, pixel_segments_raw)
    print(f"[DEBUG] Stage5转换: connect_area总数={len(connect_areas)}, 门区域={len(door_areas)}, "
          f"pixel_segments_raw={len(pixel_segments_raw)}, split_pixel_segments={len(split_pixel_segments) if split_pixel_segments else 0}, "
          f"blocks={len(raw_blocks)}")
    shapes = _build_shapes_world(raw_blocks, furniture, connect_areas, windows_from_furniture, vertices, segments, W, H, split_pixel_segments)

    output = {
        "image_size": {"width": W, "height": H},
        "floors": [{
            "id": 0,
            "name": "1楼",
            "vertex-xy": [{"id": i, "x": v[0], "y": v[1]} for i, v in enumerate(vertices)],
            "segment": segments,
            "shapes": shapes,
            "cadInfo": {
                "cadBoundingBox": {"z_max": "1.5", "z_min": "-1.5"}
            },
            "subgroup": 0,
            "tagging": [],
        }]
    }

    with open(output_path, "w", encoding="utf-8") as f:
        json.dump(output, f, indent=2, ensure_ascii=False)

    print(f"输出: {output_path}")
    print(f"  顶点数: {len(vertices)}")
    print(f"  线段数: {len(segments)}")
    print(f"  物体数: {len(shapes)} (家具 {len([s for s in shapes if not s['category'].startswith('Tag') and s['category'] not in ('SingleDoor', 'SingleWindow')])}, "
          f"房间 {len([s for s in shapes if s['category'].startswith('Tag')])}, "
          f"门 {len([s for s in shapes if s['category'] == 'SingleDoor'])}, "
          f"窗 {len([s for s in shapes if s['category'] == 'SingleWindow'])})")

    return output


# ========================================================================
# Shared pipeline runner (5 stages)
# ========================================================================

def run_5stage(work_dir, rgb_name, rgb_path, im_path, rm_path, base_name,
               initial_model, refine_model, room_model, furniture_model,
               merge_threshold, angle, dilation_kernel_size,
               center_threshold, expand_pixel, min_rect_short_side):
    """执行完整的 5 步管线，返回结果 JSON。"""
    results = {}

    # ── Stage 1: initial_mask ──────────────────────────────────
    if os.path.exists(im_path):
        print(f"[1/5] initial_mask 已存在，跳过")
    else:
        print(f"[1/5] 生成 initial_mask ...")
        generate_initial_mask(rgb_path, im_path, initial_model, rgb_name)
        print(f"  OK {im_path}")

    # ── Stage 2: refine_mask ───────────────────────────────────
    if os.path.exists(rm_path):
        print(f"[2/5] refine_mask 已存在，跳过")
    else:
        print(f"[2/5] 生成 refine_mask ...")
        generate_refine_mask(im_path, rm_path, refine_model)
        print(f"  OK {rm_path}")

    # ── Stage 3: pipeline → JSON ───────────────────────────────
    print(f"[3/5] 运行管线 ...")
    rgb_img = cv2.imread(rgb_path)
    block_mask = cv2.imread(rm_path, cv2.IMREAD_GRAYSCALE)
    if rgb_img is None:
        raise RuntimeError(f"无法读取 RGB: {rgb_path}")
    if block_mask is None:
        raise RuntimeError(f"无法读取掩码: {rm_path}")

    result_json = run_pipeline(
        rgb_img, block_mask,
        room_model, furniture_model,
        merge_threshold, angle,
        dilation_kernel_size, center_threshold,
        expand_pixel, min_rect_short_side,
    )

    out_json = os.path.join(work_dir, f"{base_name}.json")
    with open(out_json, "w", encoding="utf-8") as f:
        json.dump(result_json, f, indent=2, ensure_ascii=False)
    print(f"  OK {out_json}")

    # ── Stage 4: vis → PNG ─────────────────────────────────────
    print(f"[4/5] 生成可视化 ...")
    vis_path = os.path.join(work_dir, f"{base_name}_vis.png")
    visualize_json(result_json, vis_path, rgb_path=rgb_path, vis_door=True)
    print(f"  OK {vis_path}")

    # ── Stage 5: pixel_to_world → demo.json ────────────────────
    print(f"[5/5] 像素坐标转世界坐标 ...")
    demo_path = os.path.join(work_dir, f"{base_name}_demo.json")
    demo_data = convert_to_demo_format(result_json, demo_path)
    print(f"  OK {demo_path}")

    results["task_id"] = work_dir.split("/")[-1]
    results["output_dir"] = work_dir
    results["vis_path"] = vis_path
    results["demo_path"] = demo_path
    results["json_path"] = out_json

    return demo_data, results


# ========================================================================
# 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(50, ge=1)
    expand_pixel: int = Field(10, ge=0)
    min_rect_short_side: int = Field(30, ge=1)


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


@app.post("/process-upload")
async def process_upload(
    image: UploadFile = File(..., description="上传的 RGB 图片"),
    initial_model: str = "initial_mask.onnx",
    refine_model: str = "./black-forest-labs/FLUX.2-klein-4B",
    room_model: str = "room_cls.pt",
    furniture_model: str = "furniture_detect.onnx",
    merge_threshold: int = 15,
    angle: int = 10,
    dilation_kernel_size: int = 5,
    center_threshold: int = 50,
    expand_pixel: int = 10,
    min_rect_short_side: int = 30,
):
    """上传图片并运行完整 5 步管线。

    调用示例:
      curl -F "image=@floor.jpg" http://SERVER_IP:8070/process-upload -o result.json
    """
    task_id = uuid.uuid4().hex[:12]
    work_dir = os.path.join(tempfile.gettempdir(), f"floorplan_{task_id}")
    os.makedirs(work_dir, exist_ok=True)

    original_name = image.filename or "upload.jpg"
    safe_name = os.path.basename(original_name)
    rgb_path = os.path.join(work_dir, safe_name)
    with open(rgb_path, "wb") as f:
        content = await image.read()
        f.write(content)

    base_name = os.path.splitext(safe_name)[0]

    im_path = os.path.join(work_dir, f"initial_mask_{safe_name}")
    rm_path = os.path.join(work_dir, f"refine_mask_{safe_name}")

    try:
        demo_data, _ = run_5stage(
            work_dir, safe_name, rgb_path, im_path, rm_path, base_name,
            initial_model, refine_model, room_model, furniture_model,
            merge_threshold, angle, dilation_kernel_size,
            center_threshold, expand_pixel, min_rect_short_side,
        )
        print(f"\n===== 完成 =====")
        return JSONResponse(content=demo_data)
    except Exception as e:
        import traceback
        traceback.print_exc()
        return JSONResponse(status_code=500, content={"success": False, "error": str(e)})


@app.post("/process")
def process(req: ProcessRequest):
    """本地文件夹模式，运行完整 5 步管线。

    Returns: 直接将完整的 JSON 数据返回给客户端
    """
    folder = req.folder
    if not os.path.isdir(folder):
        return JSONResponse(status_code=400, content={"success": False, "error": f"文件夹不存在: {folder}"})

    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 JSONResponse(status_code=400, content={"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]

    out_dir = os.path.join(folder, base_name)
    os.makedirs(out_dir, exist_ok=True)

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

    im_path = os.path.join(out_dir, f"initial_mask_{rgb_name}")
    rm_path = os.path.join(out_dir, f"refine_mask_{rgb_name}")

    try:
        demo_data, _ = run_5stage(
            out_dir, rgb_name, rgb_in_out, im_path, rm_path, base_name,
            req.initial_model, req.refine_model, 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,
        )
        print(f"\n===== 完成 =====")
        return JSONResponse(content=demo_data)
    except Exception as e:
        import traceback
        traceback.print_exc()
        return JSONResponse(status_code=500, content={"success": False, "error": str(e)})


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