From 88010c5914e72d32339b30d616f3810dc2b3e861 Mon Sep 17 00:00:00 2001
From: Jim Martens <github@2martens.de>
Date: Thu, 13 Jun 2019 14:12:13 +0200
Subject: [PATCH] Added type hinting for attributes of SSD wrapper objects

Signed-off-by: Jim Martens <github@2martens.de>
---
 src/twomartens/masterthesis/ssd.py | 20 +++++++++++++++-----
 1 file changed, 15 insertions(+), 5 deletions(-)

diff --git a/src/twomartens/masterthesis/ssd.py b/src/twomartens/masterthesis/ssd.py
index 84e2429..f63f5c3 100644
--- a/src/twomartens/masterthesis/ssd.py
+++ b/src/twomartens/masterthesis/ssd.py
@@ -38,7 +38,7 @@ import math
 import os
 import pickle
 import time
-from typing import Dict, List, Sequence
+from typing import Dict, List, Sequence, Union
 from typing import Optional
 
 import numpy as np
@@ -71,6 +71,11 @@ class SSD:
     Args:
         mode: one of training, inference, and inference_fast
         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:
@@ -78,8 +83,8 @@ class SSD:
             keras_ssd300.ssd_300(image_size=IMAGE_SIZE, n_classes=N_CLASSES,
                                  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],
-                                 return_predictor_sizes=True)
-        self.mode = mode
+                                 return_predictor_sizes=True)  # type: tf.keras.models.Model, np.ndarray
+        self.mode = mode  # type: str
         
         # load existing weights
         if weights_path is not None:
@@ -108,6 +113,11 @@ class DropoutSSD:
     Args:
         mode: one of training, inference, and inference_fast
         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:
@@ -118,8 +128,8 @@ class DropoutSSD:
                                                  iou_threshold=IOU_THRESHOLD,
                                                  top_k=TOP_K,
                                                  scales=[0.07, 0.15, 0.33, 0.51, 0.69, 0.87, 1.05],
-                                                 return_predictor_sizes=True)
-        self.mode = mode
+                                                 return_predictor_sizes=True)  # type: tf.keras.models.Model, np.ndarray
+        self.mode = mode  # type: str
 
         # load existing weights
         if weights_path is not None: