Added type hinting for attributes of SSD wrapper objects

Signed-off-by: Jim Martens <github@2martens.de>
This commit is contained in:
2019-06-13 14:12:13 +02:00
parent cc9b3abe9b
commit 88010c5914

View File

@ -38,7 +38,7 @@ import math
import os import os
import pickle import pickle
import time import time
from typing import Dict, List, Sequence from typing import Dict, List, Sequence, Union
from typing import Optional from typing import Optional
import numpy as np import numpy as np
@ -71,6 +71,11 @@ class SSD:
Args: Args:
mode: one of training, inference, and inference_fast mode: one of training, inference, and inference_fast
weights_path: path to trained weights weights_path: path to trained weights
Attributes:
mode: one of training, inference, and inference_fast
predictor_sizes: sizes of predictor layers
model: Keras SSD model
""" """
def __init__(self, mode: str, weights_path: Optional[str] = None) -> None: def __init__(self, mode: str, weights_path: Optional[str] = None) -> None:
@ -78,8 +83,8 @@ class SSD:
keras_ssd300.ssd_300(image_size=IMAGE_SIZE, n_classes=N_CLASSES, keras_ssd300.ssd_300(image_size=IMAGE_SIZE, n_classes=N_CLASSES,
mode=mode, iou_threshold=IOU_THRESHOLD, top_k=TOP_K, mode=mode, iou_threshold=IOU_THRESHOLD, top_k=TOP_K,
scales=[0.07, 0.15, 0.33, 0.51, 0.69, 0.87, 1.05], scales=[0.07, 0.15, 0.33, 0.51, 0.69, 0.87, 1.05],
return_predictor_sizes=True) return_predictor_sizes=True) # type: tf.keras.models.Model, np.ndarray
self.mode = mode self.mode = mode # type: str
# load existing weights # load existing weights
if weights_path is not None: if weights_path is not None:
@ -108,6 +113,11 @@ class DropoutSSD:
Args: Args:
mode: one of training, inference, and inference_fast mode: one of training, inference, and inference_fast
weights_path: path to trained weights weights_path: path to trained weights
Attributes:
mode: one of training, inference, and inference_fast
predictor_sizes: sizes of predictor layers
model: Keras SSD model
""" """
def __init__(self, mode: str, weights_path: Optional[str] = None) -> None: def __init__(self, mode: str, weights_path: Optional[str] = None) -> None:
@ -118,8 +128,8 @@ class DropoutSSD:
iou_threshold=IOU_THRESHOLD, iou_threshold=IOU_THRESHOLD,
top_k=TOP_K, top_k=TOP_K,
scales=[0.07, 0.15, 0.33, 0.51, 0.69, 0.87, 1.05], scales=[0.07, 0.15, 0.33, 0.51, 0.69, 0.87, 1.05],
return_predictor_sizes=True) return_predictor_sizes=True) # type: tf.keras.models.Model, np.ndarray
self.mode = mode self.mode = mode # type: str
# load existing weights # load existing weights
if weights_path is not None: if weights_path is not None: