bracket_batch.py

import os
import json
import cv2
from tqdm import tqdm

# =========================
# 使用者只需設定這裡
# =========================
IMAGE_DIR = "/workspace/dataset/Tchai_4"
STAFFINFO_DIR = "/workspace/dataset/Tchai_4_staffinfo"
BRACKET_IMG_DIR = "/workspace/dataset/Tchai_4_bracket_plot"
BRACKET_JSON_DIR = "/workspace/dataset/Tchai_4_bracket"

os.makedirs(BRACKET_IMG_DIR, exist_ok=True)
os.makedirs(BRACKET_JSON_DIR, exist_ok=True)

# =========================
# 你原本就有的 function
# =========================
def find_brackets(binary_img, start_x, search_width=50, search_buffer = 10):
    """
    在 start_x 左側搜尋 brackets
    :param binary_img: 二值化後的圖片 (背景0, 前景255)
    :param start_x: 偵測到的五線譜起始 X 座標
    :param staff_groups: 之前偵測到的五線譜分組資訊 (含 y 座標範圍)
    :param search_width: 往左搜尋的寬度
    """
    # 1. 取得左側區域的 ROI
    x_min = max(0, start_x - search_width)
    roi = binary_img[:, x_min:start_x+search_buffer]
    
    # 2. 強化垂直特徵 (使用垂直 Kernel)
    kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (1, 10))
    vertical_only = cv2.morphologyEx(roi, cv2.MORPH_OPEN, kernel)
    
    # 3. 找出連通元件 (每個垂直線段)
    num_labels, labels, stats, centroids = cv2.connectedComponentsWithStats(vertical_only, connectivity=8)
    
    detected_brackets = []
    
    # 4. 根據高度篩選
    # stats 的格式: [x, y, width, height, area]
    for i in range(1, num_labels):
        x, y, w, h, area = stats[i]
        
        # 門檻值：高度必須大於一定的比例 (例如大於一個五線譜間距)
        # 這裡可以根據你的 staff_groups 資訊來動態決定
        if h > 20: 
            # 將 ROI 座標轉換回原圖座標
            global_x = x + x_min
            detected_brackets.append({
                'rect': (global_x, y, w, h),
                'center_y': y + h // 2
            })
            
    return detected_brackets

def group_brackets_to_objects(filtered_brackets, thres_1=8, thres_2=8, v_gap=5, min_h=20):
    """
    根據橫向寬度限制與垂直連續性，將碎片化的 box 組合為完整的物件。
    
    :param filtered_brackets: 原始偵測到的 box 列表 (list of dicts)
    :param thres_1: 寬度限制 (合併後的總寬度必須大於此值)
    :param thres_2: 橫向間距限制 (小於此值的 box 橫向視為連續)
    :param v_gap: 垂直容許間隔 (允許縱向上有幾像素的斷裂仍視為同一物件)
    :param min_h: 物件最終的最小高度限制
    """
    if not filtered_brackets:
        return []

    # 1. 取得 y 的總範圍
    all_y = []
    for b in filtered_brackets:
        x, y, w, h = b['rect']
        all_y.extend([y, y + h])
    y_min, y_max = int(min(all_y)), int(max(all_y))
    
    # 2. 逐行掃描 (Row-wise Analysis)
    row_segments = {} # 儲存每一行經過合併與過濾後的 x 區間
    
    for y in range(y_min, y_max + 1):
        intervals = []
        for b in filtered_brackets:
            bx, by, bw, bh = b['rect']
            if by <= y < by + bh:
                intervals.append((bx, bx + bw))
        
        if not intervals:
            continue
            
        # 橫向合併 (Horizontal Merging)
        intervals.sort()
        merged = []
        if intervals:
            curr_x1, curr_x2 = intervals[0]
            for i in range(1, len(intervals)):
                next_x1, next_x2 = intervals[i]
                if next_x1 - curr_x2 <= thres_2: # 判斷橫向間距
                    curr_x2 = max(curr_x2, next_x2)
                else:
                    merged.append((curr_x1, curr_x2))
                    curr_x1, curr_x2 = next_x1, next_x2
            merged.append((curr_x1, curr_x2))
            
        # 寬度過濾 (Width Filter - thres_1)
        valid = [seg for seg in merged if (seg[1] - seg[0]) >= thres_1]
        
        if valid:
            row_segments[y] = valid

    # 3. 縱向聚合 (Vertical Clustering)
    final_objects = [] # 儲存格式: list of [(y, x1, x2), ...]
    
    sorted_ys = sorted(row_segments.keys())
    for y in sorted_ys:
        for x1, x2 in row_segments[y]:
            found_match = False
            # 檢查是否能接到現有物件 (考量 v_gap)
            for obj in final_objects:
                last_y, last_x1, last_x2 = obj[-1]
                # 如果在垂直容許範圍內且 X 軸有重疊 (Overlap)
                if last_y < y <= last_y + v_gap:
                    if not (x2 < last_x1 or x1 > last_x2):
                        obj.append((y, x1, x2))
                        found_match = True
                        break
            
            if not found_match:
                final_objects.append([(y, x1, x2)])

    # 4. 轉換為最終 Bounding Box
    new_bracket_boxes = []
    for obj in final_objects:
        ys = [p[0] for p in obj]
        x1s = [p[1] for p in obj]
        x2s = [p[2] for p in obj]
        
        h = max(ys) - min(ys) + 1
        if h >= min_h:
            bx = min(x1s)
            by = min(ys)
            bw = max(x2s) - bx
            new_bracket_boxes.append({
                'rect': (int(bx), int(by), int(bw), int(h)),
                'center_y': int(by + h // 2)
            })
            
    return new_bracket_boxes


# =========================
# 工具函式
# =========================
def load_image(page_id):
    img_path = os.path.join(IMAGE_DIR, f"Tchai_4_{page_id}.png")
    if not os.path.exists(img_path):
        raise FileNotFoundError(img_path)
    img = cv2.imread(img_path, cv2.IMREAD_GRAYSCALE)
    return img


def load_staff_json(page_id):
    json_path = os.path.join(
        STAFFINFO_DIR, f"Tchai_4_{page_id}_yolostaff.json"
    )
    if not os.path.exists(json_path):
        raise FileNotFoundError(json_path)
    with open(json_path, "r") as f:
        return json.load(f)


def save_result_image(page_id, img):
    out_path = os.path.join(
        BRACKET_IMG_DIR, f"Tchai_4_{page_id}_bracket.png"
    )
    cv2.imwrite(out_path, img)


def save_result_json(page_id, data):
    out_path = os.path.join(
        BRACKET_JSON_DIR, f"Tchai_4_{page_id}_bracket.json"
    )
    with open(out_path, "w") as f:
        json.dump(data, f, indent=4)


# =========================
# 單頁處理主流程
# =========================
def process_page(page_id):
    img = load_image(page_id)
    staff_data = load_staff_json(page_id)

    imgh, imgw = img.shape

    start_loc = int(staff_data["avg_start_x"] * imgw)
    avg_staff_height = staff_data["staff_height"]
    staff_y_centers = staff_data["staff_y_centers"]

    # ---- bracket detection ----
    brackets = find_brackets(img, start_loc)

    # height filter（完全照你原邏輯）
    th_h_min = avg_staff_height * 2 * imgh
    th_h_max = avg_staff_height * 5 * imgh
    filtered = [
        b for b in brackets
        if th_h_min < b["rect"][3] < th_h_max
    ]
    th_w_max = 5
    filtered_brackets = [b for b in filtered if b['rect'][2] < th_w_max]
    new_boxes = group_brackets_to_objects(filtered_brackets)

    # ---- draw ----
    result_img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)
    for b in new_boxes:
        x, y, w, h = b["rect"]
        cv2.rectangle(
            result_img,
            (x, y),
            (x + w, y + h),
            (0, 255, 0),
            3
        )
    save_result_image(page_id, result_img)

    # ---- staff center → pixel ----
    staff_y_pixels = [y * imgh for y in staff_y_centers]

    # ---- bracket grouping ----
    bracket_groups = []
    sorted_boxes = sorted(new_boxes, key=lambda b: b["rect"][1])

    for b in sorted_boxes:
        x, y, w, h = b["rect"]
        y_top, y_bottom = y, y + h

        group = []
        for idx, staff_y in enumerate(staff_y_pixels):
            if (y_top - 5) <= staff_y <= (y_bottom + 5):
                group.append(idx)

        if len(group) > 1:
            bracket_groups.append(group)

    # ---- serialize boxes ----
    bracket_boxes = []
    for b in new_boxes:
        x, y, w, h = b["rect"]
        bracket_boxes.append({
            "x": int(x),
            "y": int(y),
            "w": int(w),
            "h": int(h)
        })

    # ---- update json ----
    staff_data["brackets_box"] = bracket_boxes
    staff_data["bracket_groups"] = bracket_groups

    save_result_json(page_id, staff_data)


# =========================
# Batch 主程式
# =========================
def main():
    json_files = sorted(
        f for f in os.listdir(STAFFINFO_DIR)
        if f.endswith("_yolostaff.json")
    )

    print(f"Found {len(json_files)} pages")

    for fname in tqdm(json_files):
        page_id = fname.split("_")[2]  # Tchai_4_001_yolostaff.json → 001
        try:
            process_page(page_id)
        except Exception as e:
            print(f"[ERROR] Page {page_id}: {e}")


if __name__ == "__main__":
    main()