make_large_cluster.py

import numpy as np
import matplotlib.pyplot as plt
from sklearn.mixture import GaussianMixture
from sklearn.cluster import KMeans, DBSCAN
import glob
import json
import os
from depth_flatten import flatten_bbox_depths
import time

total_bboxes = 0
answer_bboxes = 0
N = 0

def remove_outliers_iqr(data, iqr_scale=1.5):
    """IQR(Interquartile Range) 방식으로 이상치 제거"""
    q1 = np.percentile(data, 25)
    q3 = np.percentile(data, 75)
    iqr = q3 - q1
    lower_bound = q1 - iqr_scale * iqr
    upper_bound = q3 + iqr_scale * iqr
    return data[(data >= lower_bound) & (data <= upper_bound)]

def find_valid_disparities(disparity_root, bbox_root):
    """모든 disparity map과 bbox 파일을 순회하며 유효 disparity 수집"""
    disparity_files = glob.glob(os.path.join(disparity_root, "**", "defom/defom_2", "*.npy"), recursive=True)
    bbox_files = glob.glob(os.path.join(bbox_root, "**", "*.npy"), recursive=True)
    
    # 150mm, 300mm 폴더 제외
    disparity_files = [
        f for f in disparity_files
        if not (os.path.basename(os.path.dirname(os.path.dirname(f))).startswith("150mm")
                or os.path.basename(os.path.dirname(os.path.dirname(f))).startswith("300mm"))
    ]
    bbox_files = [
        f for f in bbox_files
        if not (os.path.basename(os.path.dirname(os.path.dirname(f))).startswith("150mm")
                or os.path.basename(os.path.dirname(os.path.dirname(f))).startswith("300mm"))
    ]

    print(f"총 disparity 파일 개수: {len(disparity_files)}")
    print(f"총 bbox 파일 개수: {len(bbox_files)}")

    # random shuffle
    disparity_files.sort()
    bbox_files.sort()
    indices = np.arange(len(disparity_files))
    np.random.seed(N)
    np.random.shuffle(indices)
    disparity_files = [disparity_files[i] for i in indices]
    bbox_files = [bbox_files[i] for i in indices]

    # 9:1 비율로 train/test 분리 (disparity_files와 bbox_files에서 동일하게 분리)
    split_idx = int(0.9 * len(disparity_files))

    test_disparity_files = disparity_files[split_idx:]
    test_bbox_files = bbox_files[split_idx:]

    disparity_files = disparity_files[:split_idx]
    bbox_files = bbox_files[:split_idx]

    #sort back to original order
    disparity_files.sort()
    bbox_files.sort()

    valid_disp_all = []
    for disp_path, bbox_path in zip(disparity_files, bbox_files):
        disparity_map = np.load(disp_path)
        bboxes = np.load(bbox_path)
        disparity_map = flatten_bbox_depths(disparity_map, bboxes)

        for bbox in bboxes:
            x1, y1, x2, y2 = map(int, bbox[:4])
            disp_crop = disparity_map[y1:y2, x1:x2]
            valid_disp = disp_crop[disp_crop > 0]
            filtered_disp = remove_outliers_iqr(valid_disp)
            if len(filtered_disp) > 0:
                valid_disp_all.append(filtered_disp)
    
    valid_disp_all = np.concatenate(valid_disp_all).reshape(-1, 1)
    # print(f"총 유효 disparity 개수: {len(valid_disp_all)}")
    return valid_disp_all, disparity_files, bbox_files, test_disparity_files, test_bbox_files


def disparity_to_height(valid_disp_all, fx=1050.0, baseline=0.06, camera_height=1.62, pallete_height=0.12):
    """Disparity 값을 절대 높이로 변환"""
    valid_disp_all = (fx * baseline) / valid_disp_all
    valid_disp_all = camera_height - pallete_height - valid_disp_all
    return valid_disp_all


# GMM 클러스터링
def gmm_clustering(valid_disp_all, max_clusters=3):
    """BIC 기반으로 최적의 GMM 모델 학습"""
    lowest_bic = np.inf
    best_gmm = None

    for n in range(3, max_clusters + 1):
        gmm = GaussianMixture(n_components=n, covariance_type='diag', reg_covar=1e-4)
        gmm.fit(valid_disp_all)
        bic = gmm.bic(valid_disp_all)
        if bic < lowest_bic:
            lowest_bic = bic
            best_gmm = gmm

    # 오름차순 정렬
    sorted_indices = np.argsort(best_gmm.means_.flatten())
    best_gmm.means_ = best_gmm.means_[sorted_indices]
    best_gmm.covariances_ = best_gmm.covariances_[sorted_indices]
    best_gmm.weights_ = best_gmm.weights_[sorted_indices]

    return best_gmm


# K-Means 클러스터링
def kmeans_clustering(valid_disp_all, max_clusters=3):
    """Elbow 방식으로 최적의 K 선택 후 KMeans 모델 학습"""
    inertias = []
    models = []

    for n in range(3, max_clusters + 1):
        kmeans = KMeans(n_clusters=n, random_state=0, n_init=10)
        kmeans.fit(valid_disp_all)
        inertias.append(kmeans.inertia_)
        models.append(kmeans)

    # Elbow 방식: inertia 변화율이 줄어드는 지점 찾기
    diffs = np.diff(inertias)
    optimal_k = np.argmin(np.abs(diffs)) + 1 if len(diffs) > 0 else 1

    best_kmeans = models[optimal_k - 1]

    return best_kmeans

def dbscan_clustering(valid_disp_all, eps=0.3, min_samples=50):
    """DBSCAN 클러스터링 수행"""
    db = DBSCAN(eps=eps, min_samples=min_samples).fit(valid_disp_all)
    labels = db.labels_
    n_clusters = len(set(labels)) - (1 if -1 in labels else 0)
    return db

# 각 bbox 내 disparity에 대한 군집 할당
def assign_clusters_to_bboxes(valid_disp_all, disparity_map_files, bboxes_files, model, method="GMM", fx=1050.0, baseline=0.06, camera_height=1.62, pallete_height=0.12):
    """각 bbox 내 disparity에 대해 GMM 또는 KMeans 군집 할당"""
    for disp_path, bbox_path in zip(disparity_map_files, bboxes_files):
        disparity_map = np.load(disp_path)
        bboxes = np.load(bbox_path)
        disparity_map = flatten_bbox_depths(disparity_map, bboxes)

        for i, bbox in enumerate(bboxes):
            x1, y1, x2, y2 = map(int, bbox[:4])
            disp_crop = disparity_map[y1:y2, x1:x2]
            valid_disp = disp_crop[disp_crop > 0]
            filtered_disp = remove_outliers_iqr(valid_disp)
            filtered_disp = (fx * baseline) / filtered_disp
            filtered_disp = camera_height - pallete_height - filtered_disp
            if len(filtered_disp) == 0:
                print(f"[{method}][bbox {i}] 유효 disparity 없음")
                continue

            if method == "GMM":
                cluster_labels = model.predict(filtered_disp.reshape(-1, 1))
            elif method == "KMeans":
                cluster_labels = model.predict(filtered_disp.reshape(-1, 1))
            elif method == "DBSCAN":
                cluster_labels = model.fit_predict(filtered_disp.reshape(-1, 1))
            else:
                raise ValueError("지원되지 않는 군집화 방법입니다.")

            unique, counts = np.unique(cluster_labels, return_counts=True)
            cluster_info = dict(zip(unique, counts))
            print(f"[{method}][bbox {i}] 군집 할당 결과: {cluster_info}")

# 시각화
def visualize_clusters(valid_disp_all, model, dataset, method_name="GMM"):
    """GMM, KMeans 또는 DBSCAN 결과를 시각화"""
    plt.figure(figsize=(10, 6))
    plt.hist(valid_disp_all, bins=100, density=True, alpha=0.3, color='gray', label='Disparity histogram')

    if method_name == "GMM":
        x = np.linspace(valid_disp_all.min(), valid_disp_all.max(), 1000).reshape(-1, 1)
        logprob = model.score_samples(x)
        responsibilities = model.predict_proba(x)
        pdf = np.exp(logprob)
        pdf_individual = responsibilities * pdf[:, np.newaxis]

        sorted_indices = np.argsort(model.means_.flatten())
        model.means_ = model.means_[sorted_indices]
        model.covariances_ = model.covariances_[sorted_indices]
        model.weights_ = model.weights_[sorted_indices]

        for i in range(model.n_components):
            plt.plot(x, pdf_individual[:, i], label=f'Cluster {i+1}', linewidth=2)
        plt.plot(x, pdf, 'k--', label='Mixture', linewidth=2, alpha=0.7)
        means = model.means_.flatten()
        plt.scatter(means, np.zeros_like(means), marker='*', s=250, c='red', edgecolor='k', zorder=5, label='Cluster means')

    elif method_name == "KMeans":
        centers = model.cluster_centers_.flatten()
        plt.scatter(centers, np.zeros_like(centers), marker='*', s=250, c='red', edgecolor='k', label='Cluster centers')

    elif method_name == "DBSCAN":
        labels = model.labels_
        unique_labels = set(labels)
        colors = plt.cm.tab10(np.linspace(0, 1, len(unique_labels)))
        for k, col in zip(unique_labels, colors):
            if k == -1:
                col = 'k'  # 노이즈는 검정색
            class_member_mask = (labels == k)
            xy = valid_disp_all[class_member_mask]
            plt.scatter(xy, np.zeros_like(xy), c=[col], label=f'Cluster {k}' if k != -1 else 'Noise', alpha=0.5)

    plt.title(f"1D Disparity Clustering ({method_name})")
    plt.xlabel("Absolute Height (m)")
    plt.ylabel("Probability density")
    plt.legend()
    plt.tight_layout()
    plt.show()
    plt.savefig(f"disparity_{method_name.lower()}_clustering_{dataset}.png")

# 정답 층수와 비교
def compare_with_answer(valid_disp_all, model, disparity_map_files, bboxes_files, answer_root, method="GMM"):
    global total_bboxes, answer_bboxes

    # total_bboxes = 0
    # answer_bboxes = 0
    
    for disp_path, bbox_path in zip(disparity_map_files, bboxes_files):
        layer_list = []
        # print(f"파일 처리 중: {disp_path}")
        disparity_map = np.load(disp_path)
        bboxes = np.load(bbox_path)
        # disparity_map = flatten_bbox_depths(disparity_map, bboxes)
        filename = os.path.basename(disp_path)
        answer_path = os.path.join(answer_root, filename.replace(".npy", ".txt"))
        with open(answer_path, 'r') as f:
            for line in f:
                line = line.strip()
                if not line:
                    continue
                data = json.loads(line)
                layer_list.append(data["layer"])

        for i, bbox in enumerate(bboxes):
            x1, y1, x2, y2 = map(int, bbox[:4])
            disp_crop = disparity_map[y1:y2, x1:x2]
            valid_disp = disp_crop[disp_crop > 0]
            filtered_disp = remove_outliers_iqr(valid_disp)
            filtered_disp = (1050.0 * 0.06) / filtered_disp
            filtered_disp = 1.62 - 0.12 - filtered_disp
            if len(filtered_disp) == 0:
                print(f"[{method}][bbox {i}] 유효 disparity 없음")
                continue

            if method == "GMM":
                cluster_labels = model.predict(filtered_disp.reshape(-1, 1))
            elif method == "KMeans":
                cluster_labels = model.predict(filtered_disp.reshape(-1, 1))
            elif method == "DBSCAN":
                cluster_labels = model.fit_predict(filtered_disp.reshape(-1, 1))
            else:
                raise ValueError("지원되지 않는 군집화 방법입니다.")

            unique, counts = np.unique(cluster_labels, return_counts=True)
            cluster_info = dict(zip(unique, counts))

            # 가장 많이 할당된 군집을 층수로 간주
            predicted_layer = max(cluster_info, key=cluster_info.get) + 1
            total_bboxes += 1
            if predicted_layer == layer_list[i]:
                answer_bboxes += 1

            # print(f"[{method}][bbox {i}] 군집 할당 결과: {cluster_info}, 예측 층수: {predicted_layer}, 정답 층수: {layer_list[i]}")
        

def main():
    datasets = [
        "2x2x3_random_1",
        "2x3x3_one",
        "2x3x3_one_1",
        "2x3x3_random",
        "3x3x3_one",
        "3x3x3_random",
        "2x3x4",
        # "60mm_r1_S_20251024_164240",
        ]

    valid_disp_all = []

    for dataset in datasets:
        disparity_root = f"/abr/coss32/workspace/dataset/layer_dataset/disparity_map_improve/{dataset}"
        bbox_root = f"/abr/coss32/workspace/dataset/layer_dataset/bbox/{dataset}"
        answer_root = f"/abr/coss32/workspace/dataset/layer_dataset/layers/{dataset}"
        print(f"=== 데이터셋: {dataset} ===")

        valid_disp_all, disparity_files, bbox_files, test_disparity_files, test_bbox_files = find_valid_disparities(disparity_root, bbox_root)
        valid_disp_all = disparity_to_height(valid_disp_all)

        gmm_model = gmm_clustering(valid_disp_all)
        # assign_clusters_to_bboxes(valid_disp_all, disparity_files, bbox_files, gmm_model, method="GMM")
        compare_with_answer(valid_disp_all, gmm_model, test_disparity_files, test_bbox_files, answer_root, method="GMM")
        visualize_clusters(valid_disp_all, gmm_model, dataset, method_name="GMM")

        # # K-Means
        # kmeans_model = kmeans_clustering(valid_disp_all)
        # assign_clusters_to_bboxes(valid_disp_all, disparity_files, bbox_files, kmeans_model, method="KMeans")
        # visualize_clusters(valid_disp_all, kmeans_model, method_name="KMeans")

        # # DBSCAN
        # dbscan_model = dbscan_clustering(valid_disp_all)
        # assign_clusters_to_bboxes(valid_disp_all, disparity_files, bbox_files, dbscan_model, method="DBSCAN")
        # visualize_clusters(valid_disp_all, dbscan_model, method_name="DBSCAN")

        print(f"총 bbox 개수: {total_bboxes}, 정답으로 예측한 bbox 개수: {answer_bboxes}, 정확도: {answer_bboxes/total_bboxes*100:.2f}%")

        # save results
        result_path = f"clustering_results_{dataset}.txt"
        with open(result_path, 'w') as f:
            f.write(f"총 bbox 개수: {total_bboxes}\n")
            f.write(f"정답으로 예측한 bbox 개수: {answer_bboxes}\n")
            f.write(f"정확도: {answer_bboxes/total_bboxes*100:.2f}%\n")

def total_results():
    datasets = [
        "2x2x3_random_1",
        "2x3x3_one",
        "2x3x3_one_1",
        "2x3x3_random",
        "3x3x3_one",
        "3x3x3_random",
        "2x3x4",
        # "60mm_r1_S_20251024_164240"
    ]
    all_valid_disp = []
    dataset_info = {}  # 각 dataset별 파일 리스트 저장

    # 1️⃣ 모든 dataset의 disparity 합치기
    for dataset in datasets:
        disparity_root = f"/abr/coss32/workspace/dataset/layer_dataset/disparity_map_improve/{dataset}"
        bbox_root = f"/abr/coss32/workspace/dataset/layer_dataset/bbox/{dataset}"
        # print(f"=== 데이터셋 처리 중: {dataset} ===")

        valid_disp, disparity_files, bbox_files, test_disparity_files, test_bbox_files = find_valid_disparities(disparity_root, bbox_root)
        valid_disp = disparity_to_height(valid_disp)
        all_valid_disp.append(valid_disp)

        dataset_info[dataset] = {
            "test_disparity_files": test_disparity_files,
            "test_bbox_files": test_bbox_files
        }

    # 2️⃣ 모든 disparity 합치기
    all_valid_disp = np.concatenate(all_valid_disp).reshape(-1, 1)
    # print(f"[모든 dataset 합침] 총 유효 disparity 개수: {len(all_valid_disp)}")

    # 3️⃣ GMM 클러스터링
    gmm_model = gmm_clustering(all_valid_disp)

    # 4️⃣ 각 dataset별로 결과 적용
    for dataset in datasets:
        # print(f"=== 결과 확인: {dataset} ===")
        valid_disp = np.concatenate([disparity_to_height(find_valid_disparities(
            f"/abr/coss32/workspace/dataset/layer_dataset/disparity_map_improve/{dataset}",
            f"/abr/coss32/workspace/dataset/layer_dataset/bbox/{dataset}")[0]
        )]).reshape(-1, 1)
        test_disparity_files = dataset_info[dataset]["test_disparity_files"]
        test_bbox_files = dataset_info[dataset]["test_bbox_files"]
        visualize_clusters(valid_disp, gmm_model, dataset, method_name="GMM")
        compare_with_answer(
            valid_disp,
            gmm_model,
            test_disparity_files,
            test_bbox_files,
            f"/abr/coss32/workspace/dataset/layer_dataset/layers/{dataset}",
            method="GMM"
        )
    dataset = "total"
    print(f"[{dataset}] 총 bbox 개수: {total_bboxes}, 정답으로 예측한 bbox 개수: {answer_bboxes}, 정확도: {answer_bboxes/total_bboxes*100:.2f}%")
    result_path = f"clustering_results_total_{dataset}.txt"
    with open(result_path, 'w') as f:
        f.write(f"총 bbox 개수: {total_bboxes}\n")
        f.write(f"정답으로 예측한 bbox 개수: {answer_bboxes}\n")
        f.write(f"정확도: {answer_bboxes/total_bboxes*100:.2f}%\n")
    
    return answer_bboxes / total_bboxes * 100
        
# N, 93.44%

if __name__ == "__main__":
    main()
    # total_results()
    # best_accuracy = 93.44
    # best_N = 71
    # N = 140
    # while(N < 1000):
    #     total_bboxes = 0
    #     answer_bboxes = 0
    #     N += 1
    #     print(f"=== 시드값: {N} ===")
    #     accuracy = total_results()
    #     if accuracy > best_accuracy:
    #         best_accuracy = accuracy
    #         best_N = N
    #         print(f"seed {N} 최고 정확도 갱신: {best_accuracy:.2f}%")
    #     else:
    #         print(f"seed {N} 정확도: {accuracy:.2f}% (최고 정확도: {best_accuracy:.2f}%, seed {best_N})")