Fix coco config

bendangnuksung · Oct 8, 2020 · 49ea5cd · 49ea5cd
1 parent b1c4d36
commit 49ea5cd
Showing 1 changed file with 3 additions and 4 deletions.
diff --git a/main.py b/main.py
@@ -4,7 +4,7 @@
 from tensorflow.python.saved_model import signature_constants
 from tensorflow.python.saved_model import tag_constants
 import os
-from config import mask_config
+from config import mask_config, CocoConfig
 from model import MaskRCNN
 
 sess = tf.Session()
@@ -13,8 +13,7 @@
 
 def get_config():
     if is_coco:
-        import coco
-        class InferenceConfig(coco.CocoConfig):
+        class InferenceConfig(CocoConfig):
             GPU_COUNT = 1
             IMAGES_PER_GPU = 1
 
@@ -115,4 +114,4 @@ def make_serving_ready(model_path, save_serve_path, version_number):
                      PATH_TO_SAVE_TENSORFLOW_SERVING_MODEL,
                      VERSION_NUMBER)
 
-print("COMPLETED")
+print("COMPLETED")