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')
