def __init__(self, cfg): self.cfg = cfg self.model = build_model(self.cfg) self.model.eval() self.metadata = MetadataCatalog.get(cfg.DATASETS.TEST[0]) checkpointer = DetectionCheckpointer(self.model) print('try load weights from: {}'.format(cfg.MODEL.WEIGHTS)) checkpointer.load(cfg.MODEL.WEIGHTS) self.transform_gen = T.ResizeShortestEdge( [cfg.INPUT.MIN_SIZE_TEST, cfg.INPUT.MIN_SIZE_TEST], cfg.INPUT.MAX_SIZE_TEST ) self.input_format = cfg.INPUT.FORMAT assert self.input_format in ["RGB", "BGR"], self.input_format
def build_evaluator(cls, cfg, dataset_name, output_folder=None): if output_folder is None: output_folder = os.path.join(cfg.OUTPUT_DIR, "inference") evaluator_list = [] evaluator_type = MetadataCatalog.get(dataset_name).evaluator_type if evaluator_type in ["coco", "coco_panoptic_seg"]: evaluator_list.append( COCOEvaluator(dataset_name, cfg, True, output_folder, dump=cfg.GLOBAL.DUMP_TRAIN)) elif evaluator_type == "pascal_voc": return PascalVOCDetectionEvaluator(dataset_name) if len(evaluator_list) == 0: raise NotImplementedError( "no Evaluator for the dataset {} with the type {}".format( dataset_name, evaluator_type)) elif len(evaluator_list) == 1: return evaluator_list[0] return DatasetEvaluators(evaluator_list)
from alfred.vis.image.get_dataset_label_map import coco_label_map_list import glob from models.data.custom_datasets.nuscenes import get_nuscenes_dicts from models.data.datasets import register_coco_instances categories = [ "trafficlight_red", "trafficlight_green", "trafficlight_black", "trafficlight_yellow", ] register_coco_instances( 'coco_tl', {}, './datasets/coco_tl/annotations/instances_train2017.json', './datasets/coco_tl/images') MetadataCatalog.get("coco_tl").thing_classes = categories force_color = [(0, 0, 255), (0, 255, 0), (255, 82, 12), (0, 255, 255)] class DefaultPredictor: def __init__(self, cfg): self.cfg = cfg self.model = build_model(self.cfg) self.model.eval() self.metadata = MetadataCatalog.get(cfg.DATASETS.TRAIN[0]) checkpointer = DetectionCheckpointer(self.model) print('try load weights from: {}'.format(cfg.MODEL.WEIGHTS)) checkpointer.load(cfg.MODEL.WEIGHTS)
categories = [ 'human.pedestrian.adult', 'human.pedestrian.child', 'human.pedestrian.wheelchair', 'human.pedestrian.stroller', 'human.pedestrian.personal_mobility', 'human.pedestrian.police_officer', 'human.pedestrian.construction_worker', 'vehicle.car', 'vehicle.bus.bendy', 'vehicle.bus.rigid', 'vehicle.truck', 'vehicle.construction', 'vehicle.emergency.ambulance', 'vehicle.emergency.police', 'vehicle.trailer' ] path = "datasets/nuScenes" get_dicts = lambda p=path, c=categories: get_nuscenes_dicts( path=p, version='v1.0-trainval', categories=c) DatasetCatalog.register("nusc_v1.0_trainval01", get_dicts) MetadataCatalog.get("nusc_v1.0_trainval01").thing_classes = categories class DefaultPredictor: def __init__(self, cfg): self.cfg = cfg self.model = build_model(self.cfg) self.model.eval() self.metadata = MetadataCatalog.get(cfg.DATASETS.TEST[0]) checkpointer = DetectionCheckpointer(self.model) print('try load weights from: {}'.format(cfg.MODEL.WEIGHTS)) checkpointer.load(cfg.MODEL.WEIGHTS) self.transform_gen = T.ResizeShortestEdge( [cfg.INPUT.MIN_SIZE_TEST, cfg.INPUT.MIN_SIZE_TEST],