Ejemplo n.º 1
0
    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
Ejemplo n.º 2
0
    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)
Ejemplo n.º 3
0
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],