车位角点检测代码
You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

113 lines
4.7KB

  1. """Configurate arguments."""
  2. import argparse
  3. # Threholds are collected by `collect_thresholds.py`.
  4. INPUT_IMAGE_SIZE = 512
  5. # 0: confidence, 1: point_shape, 2: offset_x, 3: offset_y, 4: cos(direction),
  6. # 5: sin(direction)
  7. NUM_FEATURE_MAP_CHANNEL = 6
  8. # image_size / 2^5 = 512 / 32 = 16
  9. FEATURE_MAP_SIZE = 80
  10. # Threshold used to filter marking points too close to image boundary
  11. BOUNDARY_THRESH = 0.05
  12. # Thresholds to determine whether an detected point match ground truth.
  13. SQUARED_DISTANCE_THRESH = 0.000277778 # 10 pixel in 600*600 image
  14. DIRECTION_ANGLE_THRESH = 0.5235987755982988 # 30 degree in rad
  15. VSLOT_MIN_DIST = 0.044771278151623496
  16. VSLOT_MAX_DIST = 0.1099427457599304
  17. HSLOT_MIN_DIST = 0.15057789144568634
  18. HSLOT_MAX_DIST = 0.44449496544202816
  19. SHORT_SEPARATOR_LENGTH = 0.199519231
  20. LONG_SEPARATOR_LENGTH = 0.46875
  21. # angle_prediction_error = 0.1384059287593468 collected from evaluate.py
  22. BRIDGE_ANGLE_DIFF = 0.09757113548987695 + 0.1384059287593468
  23. SEPARATOR_ANGLE_DIFF = 0.284967562063968 + 0.1384059287593468
  24. SLOT_SUPPRESSION_DOT_PRODUCT_THRESH = 0.8
  25. # precision = 0.995585, recall = 0.995805
  26. CONFID_THRESH_FOR_POINT = 0.11676871
  27. def add_common_arguments(parser):
  28. """Add common arguments for training and inference."""
  29. parser.add_argument('--detector_weights', default='',
  30. help="The weights of pretrained detector.")
  31. parser.add_argument('--depth_factor', type=int, default=32,
  32. help="Depth factor.")
  33. parser.add_argument('--disable_cuda', action='store_true',
  34. help="Disable CUDA.")
  35. parser.add_argument('--gpu_id', type=int, default=0,
  36. help="Select which gpu to use.")
  37. def get_parser_for_training():
  38. """Return argument parser for training."""
  39. parser = argparse.ArgumentParser()
  40. parser.add_argument('--dataset_directory', default='/home/thsw/ssd/zjc/AngularPoint1130/train',
  41. help="The location of dataset.")
  42. parser.add_argument('--cfg', type=str, default='models/yolov5s.yaml',
  43. help='model.yaml path')
  44. parser.add_argument('--optimizer_weights',
  45. help="The weights of optimizer.")
  46. parser.add_argument('--batch_size', type=int, default=56,
  47. help="Batch size.")
  48. parser.add_argument('--data_loading_workers', type=int, default=8,
  49. help="Number of workers for data loading.")
  50. parser.add_argument('--num_epochs', type=int, default=300,
  51. help="Number of epochs to train for.")
  52. parser.add_argument('--lr', type=float, default=1e-4,
  53. help="The learning rate of back propagation.")
  54. parser.add_argument('--enable_visdom', action='store_true',
  55. help="Enable Visdom to visualize training progress")
  56. parser.add_argument('--hyp', type=str, default='data/hyp.scratch.yaml',
  57. help='hyperparameters path')
  58. add_common_arguments(parser)
  59. return parser
  60. def get_parser_for_evaluation():
  61. """Return argument parser for testing."""
  62. parser = argparse.ArgumentParser()
  63. parser.add_argument('--dataset_directory', required=True,
  64. help="The location of dataset.")
  65. parser.add_argument('--enable_visdom', action='store_true',
  66. help="Enable Visdom to visualize training progress")
  67. add_common_arguments(parser)
  68. return parser
  69. def get_parser_for_ps_evaluation():
  70. """Return argument parser for testing."""
  71. parser = argparse.ArgumentParser()
  72. parser.add_argument('--label_directory', required=True,
  73. help="The location of dataset.")
  74. parser.add_argument('--image_directory', required=True,
  75. help="The location of dataset.")
  76. parser.add_argument('--enable_visdom', action='store_true',
  77. help="Enable Visdom to visualize training progress")
  78. add_common_arguments(parser)
  79. return parser
  80. def get_parser_for_inference():
  81. """Return argument parser for inference."""
  82. parser = argparse.ArgumentParser()
  83. parser.add_argument('--mode', required=True, choices=['image', 'video'],
  84. help="Inference image or video.")
  85. parser.add_argument('--video',
  86. help="Video path if you choose to inference video.")
  87. parser.add_argument('--inference_slot', action='store_true',
  88. help="Perform slot inference.")
  89. parser.add_argument('--thresh', type=float, default=0.5,
  90. help="Detection threshold.")
  91. parser.add_argument('--save', action='store_true',
  92. help="Save detection result to file.")
  93. add_common_arguments(parser)
  94. return parser