""" Tester Author: Xiaoyang Wu (xiaoyang.wu.cs@gmail.com) Please cite our work if the code is helpful to you. """ import os import time import numpy as np from collections import OrderedDict import torch import torch.distributed as dist import torch.nn.functional as F import torch.utils.data from .defaults import create_ddp_model import pointcept.utils.comm as comm from pointcept.datasets import build_dataset, collate_fn from pointcept.models import build_model from pointcept.utils.logger import get_root_logger from pointcept.utils.registry import Registry from pointcept.utils.misc import ( AverageMeter, intersection_and_union, intersection_and_union_gpu, make_dirs, ) TESTERS = Registry("testers") class TesterBase: def __init__(self, cfg, model=None, test_loader=None, verbose=False) -> None: torch.multiprocessing.set_sharing_strategy("file_system") self.logger = get_root_logger( log_file=os.path.join(cfg.save_path, "test.log"), file_mode="a" if cfg.resume else "w", ) self.logger.info("=> Loading config ...") self.cfg = cfg self.verbose = verbose if self.verbose: self.logger.info(f"Save path: {cfg.save_path}") self.logger.info(f"Config:\n{cfg.pretty_text}") if model is None: self.logger.info("=> Building model ...") self.model = self.build_model() else: self.model = model if test_loader is None: self.logger.info("=> Building test dataset & dataloader ...") self.test_loader = self.build_test_loader() else: self.test_loader = test_loader def build_model(self): model = build_model(self.cfg.model) n_parameters = sum(p.numel() for p in model.parameters() if p.requires_grad) self.logger.info(f"Num params: {n_parameters}") model = create_ddp_model( model.cuda(), broadcast_buffers=False, find_unused_parameters=self.cfg.find_unused_parameters, ) if os.path.isfile(self.cfg.weight): self.logger.info(f"Loading weight at: {self.cfg.weight}") checkpoint = torch.load(self.cfg.weight) weight = OrderedDict() for key, value in checkpoint["state_dict"].items(): if key.startswith("module."): if comm.get_world_size() == 1: key = key[7:] # module.xxx.xxx -> xxx.xxx else: if comm.get_world_size() > 1: key = "module." + key # xxx.xxx -> module.xxx.xxx weight[key] = value model.load_state_dict(weight, strict=True) self.logger.info( "=> Loaded weight '{}' (epoch {})".format( self.cfg.weight, checkpoint["epoch"] ) ) else: raise RuntimeError("=> No checkpoint found at '{}'".format(self.cfg.weight)) return model def build_test_loader(self): test_dataset = build_dataset(self.cfg.data.test) print("test_dataset: ", test_dataset.get_data(0)['coord'].shape) if comm.get_world_size() > 1: test_sampler = torch.utils.data.distributed.DistributedSampler(test_dataset) else: test_sampler = None print("test_sampler: ", test_sampler) test_loader = torch.utils.data.DataLoader( test_dataset, batch_size=self.cfg.batch_size_test_per_gpu, shuffle=False, num_workers=self.cfg.batch_size_test_per_gpu, pin_memory=True, sampler=test_sampler, collate_fn=self.__class__.collate_fn, ) def analyze_dataloader(dataloader): """ Analyze the contents and structure of a PyTorch DataLoader. Args: dataloader: PyTorch DataLoader object to analyze """ # Basic information print("DataLoader Info:") print(f"Number of batches: {len(dataloader)}") print(f"Batch size: {dataloader.batch_size}") print(f"Number of workers: {dataloader.num_workers}") print(f"Dataset type: {type(dataloader.dataset)}") # Get first batch to analyze structure first_batch = next(iter(dataloader)) if isinstance(first_batch, (tuple, list)): print("\nBatch structure (tuple/list):") for i, item in enumerate(first_batch): print(f"Item {i}:") if isinstance(item, dict): for key, value in item.items(): if isinstance(value, torch.Tensor): print(f" {key}: Tensor shape {value.shape}, dtype {value.dtype}") else: print(f" {key}: {type(value)}") elif isinstance(item, torch.Tensor): print(f" Tensor shape: {item.shape}, dtype: {item.dtype}") else: print(f" Type: {type(item)}") elif isinstance(first_batch, dict): print("\nBatch structure (dict):") for key, value in first_batch.items(): if isinstance(value, torch.Tensor): print(f" {key}: Tensor shape {value.shape}, dtype {value.dtype}") else: print(f" {key}: {type(value)}") analyze_dataloader(test_loader) # print("test_loader: ", test_loader) #check firsy batch of test_loader return test_loader def test(self): raise NotImplementedError @staticmethod def collate_fn(batch): raise collate_fn(batch) @TESTERS.register_module() class SemSegTester(TesterBase): def test(self): assert self.test_loader.batch_size == 1 logger = get_root_logger() logger.info(">>>>>>>>>>>>>>>> Start Evaluation >>>>>>>>>>>>>>>>") batch_time = AverageMeter() intersection_meter = AverageMeter() union_meter = AverageMeter() target_meter = AverageMeter() self.model.eval() save_path = os.path.join(self.cfg.save_path, "result") make_dirs(save_path) # create submit folder only on main process if ( self.cfg.data.test.type == "ScanNetDataset" or self.cfg.data.test.type == "ScanNet200Dataset" ) and comm.is_main_process(): make_dirs(os.path.join(save_path, "submit")) elif ( self.cfg.data.test.type == "SemanticKITTIDataset" and comm.is_main_process() ): make_dirs(os.path.join(save_path, "submit")) elif self.cfg.data.test.type == "NuScenesDataset" and comm.is_main_process(): import json make_dirs(os.path.join(save_path, "submit", "lidarseg", "test")) make_dirs(os.path.join(save_path, "submit", "test")) submission = dict( meta=dict( use_camera=False, use_lidar=True, use_radar=False, use_map=False, use_external=False, ) ) with open( os.path.join(save_path, "submit", "test", "submission.json"), "w" ) as f: json.dump(submission, f, indent=4) comm.synchronize() record = {} # fragment inference for idx, data_dict in enumerate(self.test_loader): end = time.time() # print("data_dict: ", data_dict) data_dict = data_dict[0] # current assume batch size is 1 # print("data_dict: ", data_dict) fragment_list = data_dict.pop("fragment_list") # print("fragment_list: ", fragment_list) segment = data_dict.pop("segment") data_name = data_dict.pop("name") pred_save_path = os.path.join(save_path, "{}_pred.npy".format(data_name)) if os.path.isfile(pred_save_path): logger.info( "{}/{}: {}, loaded pred and label.".format( idx + 1, len(self.test_loader), data_name ) ) continue #pred = np.load(pred_save_path) else: pred = torch.zeros((segment.size, self.cfg.data.num_classes)).cuda() print("fragment_coord_shapre", fragment_list[0]["coord"].shape) print("fragment_gird_coord_shape", fragment_list[0]["grid_coord"].shape) # print("fragment_list: ", len(fragment_list)) # print("fragment_list: ", fragment_list) for i in range(len(fragment_list)): # for i in range(1): fragment_batch_size = 1 s_i, e_i = i * fragment_batch_size, min( (i + 1) * fragment_batch_size, len(fragment_list) ) input_dict = collate_fn(fragment_list[s_i:e_i]) for key in input_dict.keys(): if isinstance(input_dict[key], torch.Tensor): input_dict[key] = input_dict[key].cuda(non_blocking=True) idx_part = input_dict["index"] with torch.no_grad(): pred_part = self.model(input_dict)["seg_logits"] # (n, k) pred_part = F.softmax(pred_part, -1) #pred[idx_part[bs:be], :] += pred_part[bs:be] if self.cfg.empty_cache: torch.cuda.empty_cache() bs = 0 for be in input_dict["offset"]: pred[idx_part[bs:be], :] += pred_part[bs:be] bs = be logger.info( "Test: {}/{}-{data_name}, Batch: {batch_idx}/{batch_num}".format( idx + 1, len(self.test_loader), data_name=data_name, batch_idx=i, batch_num=len(fragment_list), ) ) #save the probability of the selected class probs = pred.max(1)[0].data.cpu().numpy() pred = pred.max(1)[1].data.cpu().numpy() np.save(pred_save_path, pred) if "origin_segment" in data_dict.keys(): assert "inverse" in data_dict.keys() pred = pred[data_dict["inverse"]] probs = probs[data_dict["inverse"]] segment = data_dict["origin_segment"] intersection, union, target = intersection_and_union( pred, segment, self.cfg.data.num_classes, self.cfg.data.ignore_index ) intersection_meter.update(intersection) union_meter.update(union) target_meter.update(target) record[data_name] = dict( intersection=intersection, union=union, target=target ) mask = union != 0 iou_class = intersection / (union + 1e-10) iou = np.mean(iou_class[mask]) acc = sum(intersection) / (sum(target) + 1e-10) m_iou = np.mean(intersection_meter.sum / (union_meter.sum + 1e-10)) m_acc = np.mean(intersection_meter.sum / (target_meter.sum + 1e-10)) batch_time.update(time.time() - end) logger.info( "Test: {} [{}/{}]-{} " "Batch {batch_time.val:.3f} ({batch_time.avg:.3f}) " "Accuracy {acc:.4f} ({m_acc:.4f}) " "mIoU {iou:.4f} ({m_iou:.4f})".format( data_name, idx + 1, len(self.test_loader), segment.size, batch_time=batch_time, acc=acc, m_acc=m_acc, iou=iou, m_iou=m_iou, ) ) if ( self.cfg.data.test.type == "ScanNetDataset" or self.cfg.data.test.type == "ScanNet200Dataset" ): np.savetxt( os.path.join(save_path, "submit", "{}.txt".format(data_name)), self.test_loader.dataset.class2id[pred].reshape([-1, 1]), fmt="%d", ) elif self.cfg.data.test.type == "SemanticKITTIDataset": # 00_000000 -> 00, 000000 sequence_name, frame_name = data_name.split("_") os.makedirs( os.path.join( save_path, "submit", "sequences", sequence_name, "predictions" ), exist_ok=True, ) pred = pred.astype(np.uint32) pred = np.vectorize( self.test_loader.dataset.learning_map_inv.__getitem__ )(pred).astype(np.uint32) pred.tofile( os.path.join( save_path, "submit", "sequences", sequence_name, "predictions", f"{frame_name}.label", ) ) os.makedirs( os.path.join( save_path, "submit", "sequences", sequence_name, "probability" ), exist_ok=True, ) # save the probability of the inference class prob_save_path = os.path.join(save_path, "submit", "sequences", sequence_name, "probability", "{}_prob.npy".format(data_name)) np.save(prob_save_path, probs) elif self.cfg.data.test.type == "NuScenesDataset": np.array(pred + 1).astype(np.uint8).tofile( os.path.join( save_path, "submit", "lidarseg", "test", "{}_lidarseg.bin".format(data_name), ) ) logger.info("Syncing ...") comm.synchronize() record_sync = comm.gather(record, dst=0) if comm.is_main_process(): record = {} for _ in range(len(record_sync)): r = record_sync.pop() record.update(r) del r intersection = np.sum( [meters["intersection"] for _, meters in record.items()], axis=0 ) union = np.sum([meters["union"] for _, meters in record.items()], axis=0) target = np.sum([meters["target"] for _, meters in record.items()], axis=0) if self.cfg.data.test.type == "S3DISDataset": torch.save( dict(intersection=intersection, union=union, target=target), os.path.join(save_path, f"{self.test_loader.dataset.split}.pth"), ) iou_class = intersection / (union + 1e-10) accuracy_class = intersection / (target + 1e-10) mIoU = np.mean(iou_class) mAcc = np.mean(accuracy_class) allAcc = sum(intersection) / (sum(target) + 1e-10) logger.info( "Val result: mIoU/mAcc/allAcc {:.4f}/{:.4f}/{:.4f}".format( mIoU, mAcc, allAcc ) ) for i in range(self.cfg.data.num_classes): logger.info( "Class_{idx} - {name} Result: iou/accuracy {iou:.4f}/{accuracy:.4f}".format( idx=i, name=self.cfg.data.names[i], iou=iou_class[i], accuracy=accuracy_class[i], ) ) logger.info("<<<<<<<<<<<<<<<<< End Evaluation <<<<<<<<<<<<<<<<<") @staticmethod def collate_fn(batch): return batch @TESTERS.register_module() class ClsTester(TesterBase): def test(self): logger = get_root_logger() logger.info(">>>>>>>>>>>>>>>> Start Evaluation >>>>>>>>>>>>>>>>") batch_time = AverageMeter() intersection_meter = AverageMeter() union_meter = AverageMeter() target_meter = AverageMeter() self.model.eval() for i, input_dict in enumerate(self.test_loader): for key in input_dict.keys(): if isinstance(input_dict[key], torch.Tensor): input_dict[key] = input_dict[key].cuda(non_blocking=True) end = time.time() with torch.no_grad(): output_dict = self.model(input_dict) output = output_dict["cls_logits"] pred = output.max(1)[1] label = input_dict["category"] intersection, union, target = intersection_and_union_gpu( pred, label, self.cfg.data.num_classes, self.cfg.data.ignore_index ) if comm.get_world_size() > 1: dist.all_reduce(intersection), dist.all_reduce(union), dist.all_reduce( target ) intersection, union, target = ( intersection.cpu().numpy(), union.cpu().numpy(), target.cpu().numpy(), ) intersection_meter.update(intersection), union_meter.update( union ), target_meter.update(target) accuracy = sum(intersection_meter.val) / (sum(target_meter.val) + 1e-10) batch_time.update(time.time() - end) logger.info( "Test: [{}/{}] " "Batch {batch_time.val:.3f} ({batch_time.avg:.3f}) " "Accuracy {accuracy:.4f} ".format( i + 1, len(self.test_loader), batch_time=batch_time, accuracy=accuracy, ) ) iou_class = intersection_meter.sum / (union_meter.sum + 1e-10) accuracy_class = intersection_meter.sum / (target_meter.sum + 1e-10) mIoU = np.mean(iou_class) mAcc = np.mean(accuracy_class) allAcc = sum(intersection_meter.sum) / (sum(target_meter.sum) + 1e-10) logger.info( "Val result: mIoU/mAcc/allAcc {:.4f}/{:.4f}/{:.4f}.".format( mIoU, mAcc, allAcc ) ) for i in range(self.cfg.data.num_classes): logger.info( "Class_{idx} - {name} Result: iou/accuracy {iou:.4f}/{accuracy:.4f}".format( idx=i, name=self.cfg.data.names[i], iou=iou_class[i], accuracy=accuracy_class[i], ) ) logger.info("<<<<<<<<<<<<<<<<< End Evaluation <<<<<<<<<<<<<<<<<") @staticmethod def collate_fn(batch): return collate_fn(batch) @TESTERS.register_module() class PartSegTester(TesterBase): def test(self): test_dataset = self.test_loader.dataset logger = get_root_logger() logger.info(">>>>>>>>>>>>>>>> Start Evaluation >>>>>>>>>>>>>>>>") batch_time = AverageMeter() num_categories = len(self.test_loader.dataset.categories) iou_category, iou_count = np.zeros(num_categories), np.zeros(num_categories) self.model.eval() save_path = os.path.join( self.cfg.save_path, "result", "test_epoch{}".format(self.cfg.test_epoch) ) make_dirs(save_path) for idx in range(len(test_dataset)): end = time.time() data_name = test_dataset.get_data_name(idx) data_dict_list, label = test_dataset[idx] pred = torch.zeros((label.size, self.cfg.data.num_classes)).cuda() batch_num = int(np.ceil(len(data_dict_list) / self.cfg.batch_size_test)) for i in range(batch_num): s_i, e_i = i * self.cfg.batch_size_test, min( (i + 1) * self.cfg.batch_size_test, len(data_dict_list) ) input_dict = collate_fn(data_dict_list[s_i:e_i]) for key in input_dict.keys(): if isinstance(input_dict[key], torch.Tensor): input_dict[key] = input_dict[key].cuda(non_blocking=True) with torch.no_grad(): pred_part = self.model(input_dict)["cls_logits"] pred_part = F.softmax(pred_part, -1) if self.cfg.empty_cache: torch.cuda.empty_cache() pred_part = pred_part.reshape(-1, label.size, self.cfg.data.num_classes) pred = pred + pred_part.total(dim=0) logger.info( "Test: {} {}/{}, Batch: {batch_idx}/{batch_num}".format( data_name, idx + 1, len(test_dataset), batch_idx=i, batch_num=batch_num, ) ) pred = pred.max(1)[1].data.cpu().numpy() category_index = data_dict_list[0]["cls_token"] category = self.test_loader.dataset.categories[category_index] parts_idx = self.test_loader.dataset.category2part[category] parts_iou = np.zeros(len(parts_idx)) for j, part in enumerate(parts_idx): if (np.sum(label == part) == 0) and (np.sum(pred == part) == 0): parts_iou[j] = 1.0 else: i = (label == part) & (pred == part) u = (label == part) | (pred == part) parts_iou[j] = np.sum(i) / (np.sum(u) + 1e-10) iou_category[category_index] += parts_iou.mean() iou_count[category_index] += 1 batch_time.update(time.time() - end) logger.info( "Test: {} [{}/{}] " "Batch {batch_time.val:.3f} " "({batch_time.avg:.3f}) ".format( data_name, idx + 1, len(self.test_loader), batch_time=batch_time ) ) ins_mIoU = iou_category.sum() / (iou_count.sum() + 1e-10) cat_mIoU = (iou_category / (iou_count + 1e-10)).mean() logger.info( "Val result: ins.mIoU/cat.mIoU {:.4f}/{:.4f}.".format(ins_mIoU, cat_mIoU) ) for i in range(num_categories): logger.info( "Class_{idx}-{name} Result: iou_cat/num_sample {iou_cat:.4f}/{iou_count:.4f}".format( idx=i, name=self.test_loader.dataset.categories[i], iou_cat=iou_category[i] / (iou_count[i] + 1e-10), iou_count=int(iou_count[i]), ) ) logger.info("<<<<<<<<<<<<<<<<< End Evaluation <<<<<<<<<<<<<<<<<") @staticmethod def collate_fn(batch): return collate_fn(batch)