make_cluster.py

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


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 load_single_disparity_and_bbox(disparity_path, bbox_path):
    """단일 disparity map과 bbox 파일에서 유효 disparity 수집"""
    disparity_map = np.load(disparity_path)
    bboxes = np.load(bbox_path)
    disparity_map = flatten_bbox_depths(disparity_map, bboxes)

    valid_disp_all = []
    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_map, bboxes


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


def gmm_clustering(valid_disp_all, max_clusters=4):
    lowest_bic = np.inf
    best_gmm = None

    for n in range(1, 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

    print(f"[GMM] 최적 군집 수: {best_gmm.n_components}")
    print(f"[GMM] 군집 평균값들: {best_gmm.means_.flatten()}")
    return best_gmm


def kmeans_clustering(valid_disp_all, max_clusters=4):
    inertias, models = [], []

    for n in range(1, 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)

    diffs = np.diff(inertias)
    optimal_k = np.argmin(np.abs(diffs)) + 1 if len(diffs) > 0 else 1

    best_kmeans = models[optimal_k - 1]
    print(f"[KMeans] 최적 군집 수: {best_kmeans.n_clusters}")
    print(f"[KMeans] 군집 중심값들: {best_kmeans.cluster_centers_.flatten()}")
    return best_kmeans


def dbscan_clustering(valid_disp_all, eps=0.3, min_samples=50):
    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)
    print(f"[DBSCAN] 감지된 클러스터 개수: {n_clusters}")
    return db


def assign_clusters_to_bboxes(disparity_map, bboxes, model, method="GMM",
                              fx=1050.0, baseline=0.06, camera_height=1.62, pallete_height=0.12):
    """각 bbox 내 disparity에 대해 GMM, KMeans 또는 DBSCAN 군집 할당"""
    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 in ["GMM", "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, method_name="GMM"):
    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]

        for i in range(model.n_components):
            plt.plot(x, pdf_individual[:, i], label=f'Cluster {i+1}', linewidth=2)

        means = model.means_.flatten()
        plt.scatter(means, np.zeros_like(means), marker='*', s=250, c='red', edgecolor='k', 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'  # 노이즈는 검정색
            xy = valid_disp_all[labels == k]
            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.png")


def main():
    # ✅ 단일 파일 경로만 지정
    disparity_path = "/abr/coss32/workspace/dataset/layer_dataset/disparity_map/3x3x3_random/defom/3x3x3_random_frame_00000.npy"
    bbox_path = "/abr/coss32/workspace/dataset/layer_dataset/bbox/3x3x3_random/3x3x3_random_frame_00000.npy"

    valid_disp_all, disparity_map, bboxes = load_single_disparity_and_bbox(disparity_path, bbox_path)
    valid_disp_all = disparity_to_height(valid_disp_all)

    # 예시: DBSCAN
    dbscan_model = dbscan_clustering(valid_disp_all)
    assign_clusters_to_bboxes(disparity_map, bboxes, dbscan_model, method="DBSCAN")
    visualize_clusters(valid_disp_all, dbscan_model, method_name="DBSCAN")

if __name__ == "__main__":
    main()