def train(args, mode, _appcfg):
name = "hmd"
for subset in ["train", "val"]:
metadata = load_and_register_dataset(name, subset, _appcfg)
cfg = config.get_cfg()
cfg.merge_from_file(
"/aimldl-cod/external/detectron2/configs/COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml"
)
cfg.DATASETS.TRAIN = ("hmd_train", "hmd_val")
cfg.DATASETS.TEST = ()
cfg.DATALOADER.NUM_WORKERS = 2
cfg.MODEL.WEIGHTS = "detectron2://COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x/137849600/model_final_f10217.pkl" # initialize from model zoo
# cfg.MODEL.WEIGHTS = "/codehub/apps/detectron2/release/model_final.pth"
cfg.SOLVER.IMS_PER_BATCH = 2
cfg.SOLVER.BASE_LR = 0.00025
# cfg.SOLVER.MAX_ITER = 350000 # 300 iterations seems good enough, but you can certainly train longer
# cfg.MODEL.ROI_HEADS.BATCH_SIZE_PER_IMAGE = 128 # faster, and good enough for this toy dataset
cfg.MODEL.ROI_HEADS.BATCH_SIZE_PER_IMAGE = 512 # faster, and good enough for this toy dataset
cfg.MODEL.ROI_HEADS.NUM_CLASSES = 3 # only has one class (ballon)
os.makedirs(cfg.OUTPUT_DIR, exist_ok=True)
trainer = DefaultTrainer(cfg)
trainer.resume_or_load(resume=False)
trainer.train()
def evaluate(args, mode, _appcfg):
name = "hmd"
#uncomment if using trainer.model
# for subset in ["train", "val"]:
# metadata = load_and_register_dataset(name, subset, _appcfg)
subset = "test"
# subset = "val"
metadata = load_and_register_dataset(name, subset, _appcfg)
dataset_name = get_dataset_name(name, subset)
dataset_dicts = DatasetCatalog.get(dataset_name)
cfg = config.get_cfg()
cfg.merge_from_file(
"/codehub/external/detectron2/configs/COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml"
)
cfg.DATASETS.TRAIN = ("hmd_train", "hmd_val")
cfg.DATASETS.TEST = (dataset_name)
cfg.OUTPUT_DIR = "/codehub/apps/detectron2/release"
cfg.MODEL.WEIGHTS = os.path.join(cfg.OUTPUT_DIR, "model_final.pth")
cfg.DATALOADER.NUM_WORKERS = 2
cfg.SOLVER.IMS_PER_BATCH = 2
cfg.SOLVER.BASE_LR = 0.00025
cfg.MODEL.ROI_HEADS.BATCH_SIZE_PER_IMAGE = 128 # faster, and good enough for this toy dataset
cfg.MODEL.ROI_HEADS.NUM_CLASSES = 3
cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.7
# mapper = DatasetMapper(cfg, False)
_loader = build_detection_test_loader(cfg, dataset_name)
evaluator = COCOEvaluator(dataset_name,
cfg,
False,
output_dir=cfg.OUTPUT_DIR)
# trainer = DefaultTrainer(cfg)
# trainer.resume_or_load(resume=True)
# model = trainer.model
# predictor = DefaultPredictor(cfg)
# model = predictor.model
file_path = cfg.MODEL.WEIGHTS
model = build_model(cfg)
DetectionCheckpointer(model).load(file_path)
inference_on_dataset(model, _loader, evaluator)
# DatasetCatalog.register("balloon_" + d, lambda d=d: get_balloon_dicts("/aimldl-dat/temp/balloon/" + d))
# MetadataCatalog.get(name+"_"+subset).set(thing_classes=["balloon"])
metadata = MetadataCatalog.get(name + "_" + subset)
metadata.thing_classes = class_ids
# metadata.thing_dataset_id_to_contiguous_id = id_map
# log.info("metadata: {}".format(metadata))
return metadata
name = "hmd"
for subset in ["train", "val"]:
metadata = load_and_register_dataset(name, subset, appcfg)
cfg = config.get_cfg()
# cfg.merge_from_file("/codehub/external/detectron2/configs/COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml")
cfg.DATALOADER.NUM_WORKERS = 0
cfg = add_export_config(cfg)
cfg.merge_from_file(
model_zoo.get_config_file(
"COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml"))
cfg.DATASETS.TRAIN = ("hmd_train", )
cfg.DATASETS.TEST = ("hmd_val", )
# cfg.DATASETS.TRAIN = ("balloon_train",)
# cfg.DATASETS.TEST = ("balloon_val",)
cfg.MODEL.ROI_HEADS.NUM_CLASSES = 3
def predict(args, mode, _appcfg):
name = "hmd"
subset = "val"
BASE_IMAGE_PATH = "/aimldl-dat/samples/Trafic_Signs"
if args.path:
BASE_IMAGE_PATH = args.path
PREDICTION_OUTPUT_PATH = "/aimldl-dat/samples/Predictions"
print("BASE_IMAGE_PATH: {}".format(BASE_IMAGE_PATH))
dataset_name = get_dataset_name(name, subset)
cfg = config.get_cfg()
# cfg.merge_from_file("/aimldl-cod/external/detectron2/configs/COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml")
cfg.merge_from_file(
model_zoo.get_config_file(
"COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml"))
cfg.DATALOADER.NUM_WORKERS = 2
cfg.SOLVER.IMS_PER_BATCH = 2
cfg.SOLVER.BASE_LR = 0.00025
# cfg.SOLVER.MAX_ITER = 300 # 300 iterations seems good enough, but you can certainly train longer
cfg.MODEL.ROI_HEADS.BATCH_SIZE_PER_IMAGE = 128 # faster, and good enough for this toy dataset
cfg.MODEL.ROI_HEADS.NUM_CLASSES = 3
cfg.OUTPUT_DIR = "/codehub/apps/detectron2/release"
cfg.MODEL.WEIGHTS = os.path.join(cfg.OUTPUT_DIR, "model_final.pth")
cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.7 # set the testing threshold for this model
cfg.DATASETS.TEST = (dataset_name)
# print("cfg: {}".format(cfg.dump()))
#Predict from a directory
metadata = MetadataCatalog.get(dataset_name).set(
thing_classes=['signage', 'traffic_light', 'traffic_sign'])
# print("Metadata: {}".format(metadata))
predictor = DefaultPredictor(cfg)
for image in os.listdir(BASE_IMAGE_PATH):
image_filepath = os.path.join(BASE_IMAGE_PATH, image)
output_path = PREDICTION_OUTPUT_PATH + "/" + image
# image_filepath = "/aimldl-dat/data-gaze/AIML_Annotation/ods_job_230119/images/images-p2-050219_AT2/291018_114342_16718_zed_l_057.jpg"
# image_filepath = "/aimldl-dat/data-gaze/AIML_Annotation/ods_job_230119/images/images-p2-050219_AT2/291018_114252_16716_zed_l_099.jpg"
# print("image_filepath: {}".format(image_filepath))
im = cv2.imread(image_filepath)
outputs = predictor(im)
# visualize_predictions(im, outputs, metadata)
v = visualizer.Visualizer(im[:, :, ::-1],
metadata=metadata,
instance_mode=ColorMode.SEGMENTATION)
v = v.draw_instance_predictions(outputs["instances"].to("cpu"))
# cv2.imwrite(output_path, v.get_image()[:, :, ::-1])
cv2.imshow('', v.get_image()[:, :, ::-1])
cv2.waitKey(0)
def __init__(self, args, config):
self.arch = args.arch
self.cfg = config.get_cfg()