example-apps/PDD/pcb-defect-detection/libs/label_name_dict/label_dict.py

   1 # -*- coding: utf-8 -*-
   2 from __future__ import division, print_function, absolute_import
   3
   4 from libs.configs import cfgs
   5
   6 if cfgs.DATASET_NAME == 'ship':
   7     NAME_LABEL_MAP = {
   8         'back_ground': 0,
   9         'ship': 1
  10     }
  11 elif cfgs.DATASET_NAME == 'aeroplane':
  12     NAME_LABEL_MAP = {
  13         'back_ground': 0,
  14         'aeroplane': 1
  15     }
  16 elif cfgs.DATASET_NAME == 'WIDER':
  17     NAME_LABEL_MAP = {
  18         'back_ground': 0,
  19         'face': 1
  20     }
  21 elif cfgs.DATASET_NAME == 'jyzdata':
  22     NAME_LABEL_MAP = {
  23         'back_ground': 0,
  24         'classone': 1,
  25         'boli': 2,
  26         'dangeboli': 3,
  27         'taoci': 4,
  28         'taoci2': 4
  29     }
  30 elif cfgs.DATASET_NAME == 'icdar':
  31     NAME_LABEL_MAP = {
  32         'back_ground': 0,
  33         'text': 1
  34     }
  35 elif cfgs.DATASET_NAME.startswith('DOTA'):
  36     NAME_LABEL_MAP = {
  37         'back_ground': 0,
  38         'roundabout': 1,
  39         'tennis-court': 2,
  40         'swimming-pool': 3,
  41         'storage-tank': 4,
  42         'soccer-ball-field': 5,
  43         'small-vehicle': 6,
  44         'ship': 7,
  45         'plane': 8,
  46         'large-vehicle': 9,
  47         'helicopter': 10,
  48         'harbor': 11,
  49         'ground-track-field': 12,
  50         'bridge': 13,
  51         'basketball-court': 14,
  52         'baseball-diamond': 15
  53     }
  54 elif cfgs.DATASET_NAME == 'pascal':
  55     NAME_LABEL_MAP = {
  56         'back_ground': 0,
  57         'aeroplane': 1,
  58         'bicycle': 2,
  59         'bird': 3,
  60         'boat': 4,
  61         'bottle': 5,
  62         'bus': 6,
  63         'car': 7,
  64         'cat': 8,
  65         'chair': 9,
  66         'cow': 10,
  67         'diningtable': 11,
  68         'dog': 12,
  69         'horse': 13,
  70         'motorbike': 14,
  71         'person': 15,
  72         'pottedplant': 16,
  73         'sheep': 17,
  74         'sofa': 18,
  75         'train': 19,
  76         'tvmonitor': 20
  77     }
  78 elif cfgs.DATASET_NAME == 'pcb':
  79     NAME_LABEL_MAP = {
  80         'back_ground': 0,
  81         'missing_hole': 1,
  82         'mouse_bite': 2,
  83         'open_circuit': 3,
  84         'short': 4,
  85         'spur': 5,
  86         'spurious_copper': 6
  87     }
  88 else:
  89     assert 'please set label dict!'
  90
  91
  92 def get_label_name_map():
  93     reverse_dict = {}
  94     for name, label in NAME_LABEL_MAP.items():
  95         reverse_dict[label] = name
  96     return reverse_dict
  97
  98 LABEl_NAME_MAP = get_label_name_map()