Added iou threshold and top_k parameters to function calls

Signed-off-by: Jim Martens <github@2martens.de>
This commit is contained in:
2019-03-20 14:55:32 +01:00
parent 9a7a906acf
commit 167899b2f8

View File

@ -53,7 +53,8 @@ class SSD:
""" """
def __init__(self, mode: str, weights_path: Optional[str] = None) -> None: def __init__(self, mode: str, weights_path: Optional[str] = None) -> None:
self._model = keras_ssd300.ssd_300(image_size=IMAGE_SIZE, n_classes=N_CLASSES, mode=mode) self._model = keras_ssd300.ssd_300(image_size=IMAGE_SIZE, n_classes=N_CLASSES,
mode=mode, iou_threshold=IOU_THRESHOLD, top_k=TOP_K)
self.mode = mode self.mode = mode
# load existing weights # load existing weights
@ -75,7 +76,8 @@ class DropoutSSD:
def __init__(self, mode: str, weights_path: Optional[str] = None) -> None: def __init__(self, mode: str, weights_path: Optional[str] = None) -> None:
self._model = keras_ssd300_dropout.ssd_300_dropout(image_size=IMAGE_SIZE, n_classes=N_CLASSES, self._model = keras_ssd300_dropout.ssd_300_dropout(image_size=IMAGE_SIZE, n_classes=N_CLASSES,
dropout_rate=DROPOUT_RATE, mode=mode) dropout_rate=DROPOUT_RATE, mode=mode,
iou_threshold=IOU_THRESHOLD, top_k=TOP_K)
self.mode = mode self.mode = mode
# load existing weights # load existing weights