bracket.py

import cv2
import numpy as np
from matplotlib import pyplot as plt
import json
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 fast_group_brackets(filtered_brackets, imgh, imgw, thres_1=8, thres_2=8, v_gap=5, min_h=20):
    """
    使用 OpenCV 遮罩優化運算速度，取代逐行掃描
    """
    if not filtered_brackets:
        return []

    # 1. 建立一個與原圖一樣大的黑色遮罩 (單通道)
    mask = np.zeros((imgh, imgw), dtype=np.uint8)

    # 2. 將所有的原始 box 畫在遮罩上 (填滿白色 255)
    for b in filtered_brackets:
        x, y, w, h = b['rect']
        cv2.rectangle(mask, (x, y), (x + w, y + h), 255, -1)

    # 3. 執行橫向合併 (相當於 thres_2)
    # 使用橫向 Closing 填補間隙
    kernel_h = cv2.getStructuringElement(cv2.MORPH_RECT, (thres_2, 1))
    mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, kernel_h)

    # 4. 執行寬度過濾 (相當於 thres_1)
    # 使用橫向 Opening 移除寬度不足的雜訊
    kernel_w = cv2.getStructuringElement(cv2.MORPH_RECT, (thres_1, 1))
    mask = cv2.morphologyEx(mask, cv2.MORPH_OPEN, kernel_w)

    # 5. 執行縱向聚合 (相當於 v_gap)
    # 使用縱向 Closing 填補垂直微小斷裂
    kernel_v = cv2.getStructuringElement(cv2.MORPH_RECT, (1, v_gap))
    mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, kernel_v)

    # 6. 使用連通元件找出最後的獨立物件
    num_labels, labels, stats, centroids = cv2.connectedComponentsWithStats(mask, connectivity=8)

    new_bracket_boxes = []
    for i in range(1, num_labels):
        x, y, w, h, area = stats[i]
        if h >= min_h:
            new_bracket_boxes.append({
                'rect': (int(x), int(y), int(w), int(h)),
                'center_y': int(y + h // 2)
            })
            
    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/Tchai_4_001_bracket.png"
json_save_path = "/workspace/dataset/Tchai_4_bracket/Tchai_4_001_bracket.json"


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']

start_loc = int(most_possible_start_loc * imgw)
brackets = find_brackets(img, start_loc)

# remove brackets with height threshold
th_h_min = avg_staff_height * 2 * imgh
th_h_max = avg_staff_height * 5 * imgh
filtered_brackets = [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_brackets if b['rect'][2] < th_w_max]

new_boxes = group_brackets_to_objects(filtered_brackets)
# new_boxes = fast_group_brackets(filtered_brackets, imgh, imgw)

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