MaskClustering/make_pkl_conf.py

import numpy as np
import os
import pickle
from tqdm.auto import tqdm
from collections import defaultdict
import matplotlib.pyplot as plt
import seaborn as sns
import torch

class PredBBoxDistrPP:
    
    SCANNET_IDS = [4, 3, 6, 5, 9, 7, 8, 10, 12, 11, 14, 13, 23, 17, 18, 24, 25, 27, 28, 47, 88, 35, 36, 42, 45, 58, 49, 54, 56, 59, 60, 63, 67, 68, 102, 71, 72, 74, 81, 83, 90, 96, 122, 416, 106, 111, 117, 126, 129, 132, 155, 166, 173, 188, 300, 199, 204, 214, 219, 253, 299, 265, 273, 352, 295, 296, 301, 305, 312, 342, 358, 364, 368, 387, 395, 396, 403, 405, 414, 443, 469, 515, 744, 1157]
    
    SCANNET_LABELS = ['table', 'door', 'ceiling lamp', 'cabinet', 'blinds', 'curtain', 'chair', 'storage cabinet', 'office chair', 'bookshelf', 'whiteboard', 'window', 'box', 
                     'monitor', 'shelf', 'heater', 'kitchen cabinet', 'sofa', 'bed', 'trash can', 'book', 'plant', 'blanket', 'tv', 'computer tower', 'refrigerator', 'jacket', 
                     'sink', 'bag', 'picture', 'pillow', 'towel', 'suitcase', 'backpack', 'crate', 'keyboard', 'rack', 'toilet', 'printer', 'poster', 'painting', 'microwave', 'shoes', 
                     'socket', 'bottle', 'bucket', 'cushion', 'basket', 'shoe rack', 'telephone', 'file folder', 'laptop', 'plant pot', 'exhaust fan', 'cup', 'coat hanger', 'light switch', 
                     'speaker', 'table lamp', 'kettle', 'smoke detector', 'container', 'power strip', 'slippers', 'paper bag', 'mouse', 'cutting board', 'toilet paper', 'paper towel', 
                     'pot', 'clock', 'pan', 'tap', 'jar', 'soap dispenser', 'binder', 'bowl', 'tissue box', 'whiteboard eraser', 'toilet brush', 'spray bottle', 'headphones', 'stapler', 'marker']
    
    ID2LABEL = dict(zip(SCANNET_IDS, SCANNET_LABELS))
    
    LABEL2ID = dict(zip(SCANNET_LABELS, SCANNET_IDS))
    
    INV_SCANNET_IDS = {idx: i for i, idx in enumerate(SCANNET_IDS)}
    
    @staticmethod
    def _normalize_scene_id(value):
        base = os.path.basename(value)
        if base.endswith('.bin'):
            base = base[:-4]
        else:
            base = os.path.splitext(base)[0]
        return base
    
    def __init__(self, path, bins_path):
        self.path = path
        self.bins_path = bins_path
        self.get_scenes()
        self.class_scores = defaultdict(list)
        for scene_id in self.scene_ids:
            self.get_scene_inst(scene_id)
        # сортировка по убыванию числа предсказаний на класс
        self.sorted_names = sorted(self.SCANNET_LABELS, key=lambda x: -len(self.class_scores.get(x, [])))
        
    # GT-PKL больше не используется
    
    def get_scenes(self):
        self.scene_ids = []
        if not os.path.isdir(self.path):
            return
        pred_files = [f for f in os.listdir(self.path) if f.endswith('.npz')]
        for fname in pred_files:
            scene_name = os.path.splitext(fname)[0]
            bin_path = os.path.join(self.bins_path, f"{scene_name}.bin")
            if os.path.exists(bin_path):
                self.scene_ids.append(scene_name)
    
    def get_scene_inst(self, scene_id):
        cls_path = f'{self.path}/{scene_id}.npz'
        cls_data = np.load(cls_path, allow_pickle=True)
        for class_id, class_score in zip(cls_data['pred_classes'], cls_data['pred_score']):
            self.class_scores[self.ID2LABEL[class_id]].append(class_score)
    
    def plot_class_distr(self, class_name='all'):
        """
        Построить распределение оценок для конкретного класса или всех классов вместе
        
        Parameters:
        class_name: str or list - название класса, 'all' для всех классов, 
                   или список названий классов
        """
        if class_name == 'all':
            # Собираем все оценки из всех классов
            all_scores = []
            for scores in self.class_scores.values():
                all_scores.extend(scores)
            scores = all_scores
            display_name = 'All Classes'
        elif isinstance(class_name, list):
            # Собираем оценки из указанных классов
            selected_scores = []
            for cls in class_name:
                if cls in self.class_scores:
                    selected_scores.extend(self.class_scores[cls])
                else:
                    print(f"Warning: Class '{cls}' not found in class_scores")
            scores = selected_scores
            display_name = f'Classes: {", ".join(class_name[:3])}{"..." if len(class_name) > 3 else ""}'
        else:
            # Один конкретный класс
            if class_name not in self.class_scores:
                print(f"Class '{class_name}' not found in class_scores")
                print(f"Available classes: {list(self.class_scores.keys())[:10]}...")
                return
            scores = self.class_scores[class_name]
            display_name = class_name
        
        if not scores:
            print(f"No scores available for: {display_name}")
            return
        
        # Создаем фигуру
        fig, ax = plt.subplots(figsize=(12, 8))
        
        # Гистограмма с KDE (seaborn) с нормализованной осью Y
        sns.histplot(scores, bins=30, kde=True, ax=ax, color='skyblue', 
                    stat='density', alpha=0.7)
        ax.set_title(f'Distribution of scores for {display_name}', fontsize=14, fontweight='bold')
        ax.set_xlabel('Score', fontsize=12)
        ax.set_ylabel('Density', fontsize=12)
        ax.grid(True, alpha=0.3)
        
        # Добавляем вертикальную линию для среднего значения
#         mean_score = np.mean(scores)
#         ax.axvline(mean_score, color='red', linestyle='--', linewidth=2, 
#                   label=f'Mean: {mean_score:.3f}')
        
        # Добавляем вертикальную линию для медианы
        median_score = np.median(scores)
        ax.axvline(median_score, color='green', linestyle='--', linewidth=2, 
                  label=f'Median: {median_score:.3f}')
        ax.axvline(np.percentile(scores, 32.45), color='red', linestyle='-', linewidth=2, 
                  label=f'Size bound: {np.percentile(scores, 32.45):.3f}')
        # Добавляем легенду
        ax.legend()
        
        # Добавляем статистику в текстовом блоке
        if class_name == 'all':
            class_info = f"Total classes: {len(self.class_scores)}"
        elif isinstance(class_name, list):
            class_info = f"Selected classes: {len(class_name)}"
        else:
            class_info = f"Class: {class_name}"
            
        stats_text = f"""Statistics for {display_name}:
        {class_info}
        Total instances: {len(scores):,}
        Mean: {np.mean(scores):.3f}
        Median: {np.median(scores):.3f}
        Std: {np.std(scores):.3f}
        Min: {np.min(scores):.3f}
        Max: {np.max(scores):.3f}
        Q1: {np.percentile(scores, 25):.3f}
        Q : {np.percentile(scores, 32.45):.3f}
        Q3: {np.percentile(scores, 75):.3f}"""
        
        # Размещаем текстовый блок в удобном месте
        props = dict(boxstyle="round,pad=0.5", facecolor="lightgray", alpha=0.8)
        ax.text(0.02, 0.98, stats_text, transform=ax.transAxes, fontfamily='monospace',
                verticalalignment='top', bbox=props, fontsize=10)
        
        plt.tight_layout()
        plt.show()
        
        # Также выводим статистику в консоль
        print(stats_text)
        
        return scores  # Возвращаем массив оценок для дальнейшего анализа

    # Дополнительный метод для сравнения нескольких классов
    def plot_multiple_classes(self, class_names: list):
        """
        Сравнить распределения нескольких классов на одном графике
        """
        fig, ax = plt.subplots(figsize=(12, 8))
        
        colors = ['skyblue', 'lightcoral', 'lightgreen', 'gold', 'lightpink']
        
        for i, cls in enumerate(class_names):
            if cls not in self.class_scores:
                print(f"Warning: Class '{cls}' not found, skipping")
                continue
                
            scores = self.class_scores[cls]
            if scores:
                sns.kdeplot(scores, ax=ax, label=cls, color=colors[i % len(colors)], 
                           linewidth=2, alpha=0.8)
        
        ax.set_title('Score Distribution Comparison', fontsize=14, fontweight='bold')
        ax.set_xlabel('Score', fontsize=12)
        ax.set_ylabel('Density', fontsize=12)
        ax.grid(True, alpha=0.3)
        ax.legend()
        
        plt.tight_layout()
        plt.show()
    
    def get_class_lowerbound(self, class_name='all', percentile=32.45):
        if class_name == 'all':
            # Собираем все оценки из всех классов
            all_scores = []
            for scores in self.class_scores.values():
                all_scores.extend(scores)
            scores = all_scores
        elif isinstance(class_name, list):
            selected_scores = []
            for cls in class_name:
                if cls in self.class_scores:
                    selected_scores.extend(self.class_scores[cls])
                else:
                    print(f"Warning: Class '{cls}' not found in class_scores")
            scores = selected_scores
        else:
            # Один конкретный класс
            if class_name not in self.class_scores:
                print(f"Class '{class_name}' not found in class_scores")
                print(f"Available classes: {list(self.class_scores.keys())[:10]}...")
                return
            scores = self.class_scores[class_name]
        
        return np.percentile(scores, percentile)
    
    def get_bboxes_by_masks(self, masks, points):
        boxes = []
        for mask in masks:
            object_points = points[mask][:, :3]
            # xyz_min = object_points.min(dim=0).values
            # xyz_max = object_points.max(dim=0).values
            xyz_min = object_points.quantile(0.01, dim=0)
            xyz_max = object_points.quantile(0.99, dim=0)
            center = (xyz_max + xyz_min) / 2
            size = xyz_max - xyz_min
            box = torch.cat((center, size))
            boxes.append(box)
        assert len(boxes) != 0, "Why 0 masks in scene?"
        boxes = torch.stack(boxes)
        return boxes
    
    def filter_instances_by_confidence(self, scene_name, threshold=0.5, class_agnostic=False):
        instances = []
        points_path = f'{self.bins_path}/{scene_name}.bin'
        if not os.path.exists(points_path):
            return instances
        points = torch.from_numpy(np.fromfile(points_path, dtype=np.float32).reshape((-1, 6)))
        cls_path = f'{self.path}/{scene_name}.npz'
        if not os.path.exists(cls_path):
            return instances
        cls_data = np.load(cls_path, allow_pickle=True)
        pred_masks = torch.from_numpy(cls_data['pred_masks']).T
        pred_classes = cls_data['pred_classes']
        pred_scores = cls_data['pred_score']
        if len(pred_scores) == 0:
            return instances
        np_indices = np.where(pred_scores >= threshold)[0]
        if np_indices.size == 0:
            return instances
        torch_indices = torch.as_tensor(np_indices, dtype=torch.long)
        selected_masks = pred_masks[torch_indices]
        boxes = self.get_bboxes_by_masks(selected_masks, points)
        selected_classes = pred_classes[np_indices]
        for box, pred_class in zip(boxes, selected_classes):
            write_class = 0 if class_agnostic else self.INV_SCANNET_IDS[pred_class]
            instances.append({'bbox_3d': box.numpy().tolist(), 'bbox_label_3d': write_class})
        return instances
    
    # Фильтрация по GT удалена; используйте filter_instances_by_confidence
    
    
    def build_pkl_by_confidence(self, pkl_path, threshold=0.5, class_agnostic=False):
        new_data: dict = {"metainfo": {'categories': {'table': 0, 'door': 1, 'ceiling lamp': 2, 'cabinet': 3, 'blinds': 4, 'curtain': 5, 'chair': 6, 'storage cabinet': 7, 'office chair': 8, 'bookshelf': 9, 'whiteboard': 10, 'window': 11, 'box': 12, 'monitor': 13, 'shelf': 14, 'heater': 15, 'kitchen cabinet': 16, 'sofa': 17, 'bed': 18, 'trash can': 19, 'book': 20, 'plant': 21, 'blanket': 22, 'tv': 23, 'computer tower': 24, 'refrigerator': 25, 'jacket': 26, 'sink': 27, 'bag': 28, 'picture': 29, 'pillow': 30, 'towel': 31, 'suitcase': 32, 'backpack': 33, 'crate': 34, 'keyboard': 35, 'rack': 36, 'toilet': 37, 'printer': 38, 'poster': 39, 'painting': 40, 'microwave': 41, 'shoes': 42, 'socket': 43, 'bottle': 44, 'bucket': 45, 'cushion': 46, 'basket': 47, 'shoe rack': 48, 'telephone': 49, 'file folder': 50, 'laptop': 51, 'plant pot': 52, 'exhaust fan': 53, 'cup': 54, 'coat hanger': 55, 'light switch': 56, 'speaker': 57, 'table lamp': 58, 'kettle': 59, 'smoke detector': 60, 'container': 61, 'power strip': 62, 'slippers': 63, 'paper bag': 64, 'mouse': 65, 'cutting board': 66, 'toilet paper': 67, 'paper towel': 68, 'pot': 69, 'clock': 70, 'pan': 71, 'tap': 72, 'jar': 73, 'soap dispenser': 74, 'binder': 75, 'bowl': 76, 'tissue box': 77, 'whiteboard eraser': 78, 'toilet brush': 79, 'spray bottle': 80, 'headphones': 81, 'stapler': 82, 'marker': 83}, 'dataset': 'scannetpp', 'info_version': '1.0'}}
        data_list = []
        for scene_name in tqdm(self.scene_ids):
            instances = self.filter_instances_by_confidence(scene_name, threshold=threshold, class_agnostic=class_agnostic)
            scene_entry = {
                'lidar_points': {
                    'num_pts_feats': 6,
                    'lidar_path': f'{scene_name}.bin',
                },
                'instances': instances,
                'pts_semantic_mask_path': f'{scene_name}.bin',
                'pts_instance_mask_path': f'{scene_name}.bin',
                'axis_align_matrix': np.eye(4, dtype=np.float32),
            }
            data_list.append(scene_entry)
        new_data['data_list'] = data_list
        with open(pkl_path, 'wb') as file:
            pickle.dump(new_data, file)
        
    @property
    def scores(self):
        return self.class_scores
            

if __name__ == "__main__":
    pred_path = \
        "/home/jovyan/users/bulat/workspace/3drec/Indoor/MaskClustering/data/prediction/scannetpp_v2_dust3r_unposed"
    bins_path = \
        "/home/jovyan/users/bulat/workspace/3drec/Indoor/OKNO/data/scannetpp/bins/points_dust3r_v2_unposed"
    out_pkl_path = \
        "/home/jovyan/users/bulat/workspace/3drec/Indoor/OKNO/data/scannetpp/bins/scannetpp84_v2_dust3r_unposed_train.pkl"
    threshold = 0.5

    distr = PredBBoxDistrPP(pred_path, bins_path)
    distr.build_pkl_by_confidence(out_pkl_path, threshold=threshold, class_agnostic=False)