pcb defect detetcion application

[ealt-edge.git] / example-apps / PDD / pcb-defect-detection / libs / label_name_dict / label_dict.py

diff --git a/example-apps/PDD/pcb-defect-detection/libs/label_name_dict/label_dict.py b/example-apps/PDD/pcb-defect-detection/libs/label_name_dict/label_dict.py
new file mode 100755 (executable)
index 0000000..dc5d504
--- /dev/null
+++ b/example-apps/PDD/pcb-defect-detection/libs/label_name_dict/label_dict.py
@@ -0,0 +1,98 @@
+# -*- coding: utf-8 -*-
+from __future__ import division, print_function, absolute_import
+
+from libs.configs import cfgs
+
+if cfgs.DATASET_NAME == 'ship':
+    NAME_LABEL_MAP = {
+        'back_ground': 0,
+        'ship': 1
+    }
+elif cfgs.DATASET_NAME == 'aeroplane':
+    NAME_LABEL_MAP = {
+        'back_ground': 0,
+        'aeroplane': 1
+    }
+elif cfgs.DATASET_NAME == 'WIDER':
+    NAME_LABEL_MAP = {
+        'back_ground': 0,
+        'face': 1
+    }
+elif cfgs.DATASET_NAME == 'jyzdata':
+    NAME_LABEL_MAP = {
+        'back_ground': 0,
+        'classone': 1,
+        'boli': 2,
+        'dangeboli': 3,
+        'taoci': 4,
+        'taoci2': 4
+    }
+elif cfgs.DATASET_NAME == 'icdar':
+    NAME_LABEL_MAP = {
+        'back_ground': 0,
+        'text': 1
+    }
+elif cfgs.DATASET_NAME.startswith('DOTA'):
+    NAME_LABEL_MAP = {
+        'back_ground': 0,
+        'roundabout': 1,
+        'tennis-court': 2,
+        'swimming-pool': 3,
+        'storage-tank': 4,
+        'soccer-ball-field': 5,
+        'small-vehicle': 6,
+        'ship': 7,
+        'plane': 8,
+        'large-vehicle': 9,
+        'helicopter': 10,
+        'harbor': 11,
+        'ground-track-field': 12,
+        'bridge': 13,
+        'basketball-court': 14,
+        'baseball-diamond': 15
+    }
+elif cfgs.DATASET_NAME == 'pascal':
+    NAME_LABEL_MAP = {
+        'back_ground': 0,
+        'aeroplane': 1,
+        'bicycle': 2,
+        'bird': 3,
+        'boat': 4,
+        'bottle': 5,
+        'bus': 6,
+        'car': 7,
+        'cat': 8,
+        'chair': 9,
+        'cow': 10,
+        'diningtable': 11,
+        'dog': 12,
+        'horse': 13,
+        'motorbike': 14,
+        'person': 15,
+        'pottedplant': 16,
+        'sheep': 17,
+        'sofa': 18,
+        'train': 19,
+        'tvmonitor': 20
+    }
+elif cfgs.DATASET_NAME == 'pcb':
+    NAME_LABEL_MAP = {
+        'back_ground': 0,
+        'missing_hole': 1,
+        'mouse_bite': 2,
+        'open_circuit': 3,
+        'short': 4,
+        'spur': 5,
+        'spurious_copper': 6
+    }
+else:
+    assert 'please set label dict!'
+
+
+def get_label_name_map():
+    reverse_dict = {}
+    for name, label in NAME_LABEL_MAP.items():
+        reverse_dict[label] = name
+    return reverse_dict
+
+LABEl_NAME_MAP = get_label_name_map()
\ No newline at end of file