bracket_v2.py

import cv2
import numpy as np
from matplotlib import pyplot as plt
import json
def refine_barlines(barlines):
    """
    融合重疊或相近的 barlines
    規則：
      - 若兩個 barline 在 y 方向重疊，且 x 距離 < 4，則合併
      - 合併後取：
          x = 最左邊
          y = 最上邊
          w = 1
          h = 最下邊 - 最上邊
    回傳：
      refine_barline (list of tuples): [(x, y, w, h), ...]
    """
    # 依 x 排序（方便鄰近比對）
    barlines = sorted(barlines, key=lambda b: b[0])
    merged = []
    used = [False] * len(barlines)

    for i, (x1, y1, w1, h1) in enumerate(barlines):
        if used[i]:
            continue
        x1_right = x1 + w1
        y1_bottom = y1 + h1
        
        # 初始化融合範圍
        new_x = x1
        new_y_top = y1
        new_y_bottom = y1_bottom
        used[i] = True

        for j in range(i + 1, len(barlines)):
            if used[j]:
                continue
            x2, y2, w2, h2 = barlines[j]
            x2_right = x2 + w2
            y2_bottom = y2 + h2

            # 判斷 y 是否重疊
            overlap_y = not (y2_bottom < new_y_top or y2 > new_y_bottom)
            # 判斷 x 是否夠近
            if abs(x2 - new_x) < 4 and overlap_y:
                # 更新上下範圍
                new_y_top = min(new_y_top, y2)
                new_y_bottom = max(new_y_bottom, y2_bottom)
                new_x = min(new_x, x2)
                used[j] = True

        merged.append((new_x, new_y_top, 1, new_y_bottom - new_y_top))

    return merged

def get_vertical_lines(image_path, debug=False):
    # 讀取灰階圖
    img = cv2.imread(image_path, cv2.IMREAD_GRAYSCALE)
    if img is None:
        raise FileNotFoundError(f"Image not found: {image_path}")

    # binary
    binary = cv2.adaptiveThreshold(
        ~img, 255,
        cv2.ADAPTIVE_THRESH_MEAN_C,
        cv2.THRESH_BINARY,
        15, -2
    )

    # 提取垂直線 (Morphology)
    vertical_kernel_len = img.shape[0] // 40  # 可調整
    print("Vertical kernel length:", vertical_kernel_len)
    vertical_kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (1, vertical_kernel_len))
    print("Vertical kernel shape:", vertical_kernel.shape) # 內部所有元素都是1
    # print("Vertical kernel:\n", vertical_kernel)
    vertical_lines = cv2.morphologyEx(binary, cv2.MORPH_OPEN, vertical_kernel, iterations=1)


    # 去掉太長的直線（例如頁邊）
    num_labels, labels, stats, _ = cv2.connectedComponentsWithStats(vertical_lines, connectivity=8)
    mask = np.zeros_like(vertical_lines)
    for i in range(1, num_labels):
        x, y, w, h, area = stats[i]
        if 30 < h < img.shape[0] * 0.9:
            mask[labels == i] = 255
    vertical_lines = mask

    # ================================
    # Step 3️⃣：分群相近 x 座標的 barlines
    # ================================
    num_labels, labels, stats, _ = cv2.connectedComponentsWithStats(vertical_lines, connectivity=8)
    barlines = []
    for i in range(1, num_labels):
        x, y, w, h, area = stats[i]
        if h > 50:
            barlines.append((x, y, w, h))
    if len(barlines) == 0:
        raise ValueError("No barlines detected. Try adjusting vertical_kernel_len or threshold.")
    print(f"Detected {len(barlines)} barlines.")  
    
    # barlines 融合
    refine_barline = refine_barlines(barlines)
    print("The number of refined barlines:", len(refine_barline))
    # remove w
    refine_barline_line = [(x, y, 1, h) for (x, y, w, h) in refine_barline]
    return refine_barline_line

def group_brackets_to_objects(filtered_brackets, thres_1=3, thres_2=10, 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
        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
            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(
(int(bx), int(by), int(bw), int(h))
            )
            
    return new_bracket_boxes

image_path = "/workspace/25-omr/orchestra_dataset/Tchai_4/Tchai_4_1.png"
json_path = "/workspace/dataset/Tchai_4_staffinfo/Tchai_4_001_yolostaff.json"
plt_save_path = "/workspace/dataset/Tchai_4_bracket_plot_v2/Tchai_4_001_bracket.png"
json_save_path = "/workspace/dataset/Tchai_4_bracket_v2/Tchai_4_001_bracket.json"


vertical_lines = get_vertical_lines(image_path, debug=True)

# plot it on the image
img = cv2.imread(image_path, cv2.IMREAD_GRAYSCALE)
imgh, imgw = img.shape

with open(json_path, "r") as f:
    staff_data = json.load(f)
most_possible_start_loc = staff_data['avg_start_x']
avg_staff_height = staff_data['staff_height']
staff_y_centers = staff_data['staff_y_centers']

# remove brackets with height threshold
start_loc = int(most_possible_start_loc * imgw)
# start_loc = 167
print(start_loc)
# for b in vertical_lines:
#     print(b[0])
filtered_brackets = [b for b in vertical_lines if start_loc-50 < b[0] < start_loc+10]
print(len(filtered_brackets))
for b in filtered_brackets:
    print(b)
print(imgw, imgh)
th_h_min = avg_staff_height * 2 * imgh
th_h_max = avg_staff_height * 10 * imgh
print(th_h_min)
print(th_h_max)
filtered_brackets = [b for b in filtered_brackets if th_h_min < b[3] < th_h_max]

new_boxes = group_brackets_to_objects(filtered_brackets)


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)

# 修正：直接使用 cv2.imwrite 存檔以保持原始大小
cv2.imwrite(plt_save_path, result_img)
print(f"存檔完成：{plt_save_path}，大小為 {result_img.shape[1]}x{result_img.shape[0]}")


# 1. 將 staff_y_centers 轉換為像素座標 (假設 imgh 是你的圖片高度)
staff_y_pixels = [y_norm * imgh for y_norm in staff_y_centers]

# 2. 建立關聯 list
bracket_groups = []

# 為了確保輸出的順序是從上到下，我們先對 new_boxes 按 y 座標排序
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
    y_bottom = y + h
    
    # 找出所有中心點落在這個 box 垂直範圍內的 staff 索引
    current_group = []
    for idx, staff_y in enumerate(staff_y_pixels):
        # 判斷 staff 中心是否在 bracket 的高度範圍內
        # 這裡可以稍微放寬 5-10 像素的容錯
        if (y_top - 5) <= staff_y <= (y_bottom + 5):
            current_group.append(idx)
    
    # 如果這個 bracket 有包住任何 staff，且數量大於 1 (通常 bracket 是為了分組)
    # 如果你連單個的也要，就把 len(current_group) > 1 拿掉
    if len(current_group) > 1:
        bracket_groups.append(current_group)

print("Final Bracket Grouping:", bracket_groups)

# store the bracket data
bracket_data = []
for b in new_boxes:
    x, y, w, h = b['rect']
    bracket_data.append({
        'x': int(x),
        'y': int(y),
        'w': int(w),
        'h': int(h)
    })

# update staff_data
staff_data['brackets_box'] = bracket_data
staff_data['bracket_groups'] = bracket_groups

with open(json_save_path, "w") as f:
    json.dump(staff_data, f, indent=4)
print(f"Bracket data saved to {json_save_path}")