def od_detection_mask_learner(od_detection_mask_dataset): """ returns a mask detection learner that has been trained for one epoch. """ model = get_pretrained_maskrcnn( num_classes=len(od_detection_mask_dataset.labels) + 1, min_size=100, max_size=200, rpn_pre_nms_top_n_train=500, rpn_pre_nms_top_n_test=250, rpn_post_nms_top_n_train=500, rpn_post_nms_top_n_test=250, ) learner = DetectionLearner(od_detection_mask_dataset, model=model) learner.fit(1) return learner
def od_detection_keypoint_learner(tiny_od_detection_keypoint_dataset): """ returns a keypoint detection learner that has been trained for one epoch. """ model = get_pretrained_keypointrcnn( num_classes=len(tiny_od_detection_keypoint_dataset.labels) + 1, num_keypoints=len( tiny_od_detection_keypoint_dataset.keypoint_meta["labels"]), min_size=100, max_size=200, rpn_pre_nms_top_n_train=500, rpn_pre_nms_top_n_test=250, rpn_post_nms_top_n_train=500, rpn_post_nms_top_n_test=250, ) learner = DetectionLearner(tiny_od_detection_keypoint_dataset, model=model) learner.fit(1, skip_evaluation=True) return learner
def test_detection_learner_init_model(od_detection_dataset): """ Tests detection learner with model settings. """ classes = len(od_detection_dataset.labels) model = get_pretrained_fasterrcnn(num_classes=classes, min_size=600, max_size=2000) learner = DetectionLearner(od_detection_dataset, model=model) assert type(learner) == DetectionLearner assert learner.model == model assert learner.model != get_pretrained_fasterrcnn(classes)
def od_detections(od_detection_dataset): """ returns output of the object detector for a given test set. """ learner = DetectionLearner(od_detection_dataset) return learner.predict_dl(od_detection_dataset.test_dl, threshold=0)
def test_detection_learner_init(od_detection_dataset): """ Tests detection learner basic init. """ learner = DetectionLearner(od_detection_dataset) assert type(learner) == DetectionLearner
def test_detection_init_from_saved_model(saved_model): """ Test that we can create an detection learner from a saved model. """ name, path = saved_model DetectionLearner.from_saved_model(name, path)
# ---------------------------------------------------------------------- # Prepare processing function # ---------------------------------------------------------------------- # Load ResNet model. model = torchvision.models.detection.fasterrcnn_resnet50_fpn( pretrained=True, rpn_pre_nms_top_n_test=5, rpn_post_nms_top_n_test=5, max_size=200, ) detector = DetectionLearner( model=model, labels=coco_labels()[1:], # First element is '__background__' ) if len(args.path): for path in args.path: if is_url(path): tempdir = tempfile.gettempdir() imfile = os.path.join(tempdir, "temp.jpg") urllib.request.urlretrieve(path, imfile) else: imfile = os.path.join(get_cmd_cwd(), path) try: im = Image.open(imfile).convert('RGB')