updates

2020-02-19 17:08:03 -08:00
parent f92ad043bd
commit 6fbab656c8
2 changed files with 23 additions and 18 deletions
@@ -25,10 +25,16 @@ def parse_model_cfg(path):
            key, val = line.split("=")
            key = key.rstrip()

-            if 'anchors' in key:
+            if key == 'anchors':  # return nparray
                mdefs[-1][key] = np.array([float(x) for x in val.split(',')]).reshape((-1, 2))  # np anchors
+            elif key in ['from', 'layers', 'mask']:  # return array
+                mdefs[-1][key] = [int(x) for x in val.split(',')]
            else:
-                mdefs[-1][key] = val.strip()
+                val = val.strip()
+                if val.isnumeric():  # return int or float
+                    mdefs[-1][key] = int(val) if (int(val) - float(val)) == 0 else float(val)
+                else:
+                    mdefs[-1][key] = val  # return string

    # Check all fields are supported
    supported = ['type', 'batch_normalize', 'filters', 'size', 'stride', 'pad', 'activation', 'layers', 'groups',