2023-08-17 11:59:31 +08:00
|
|
|
"""Configurate arguments."""
|
|
|
|
|
import argparse
|
|
|
|
|
|
|
|
|
|
# 0: confidence, 1: point_shape, 2: offset_x, 3: offset_y, 4: cos(direction),
|
|
|
|
|
# 5: sin(direction)
|
|
|
|
|
NUM_FEATURE_MAP_CHANNEL = 6
|
|
|
|
|
|
|
|
|
|
def add_common_arguments(parser):
|
|
|
|
|
"""Add common arguments for training and inference."""
|
|
|
|
|
parser.add_argument('--detector_weights', default=r'E:\pycharmProject\DMPR-PS\weights\dp_detector_499.pth',
|
|
|
|
|
help="The weights of pretrained detector.")
|
2023-09-19 16:02:42 +08:00
|
|
|
parser.add_argument('--cfg', type=str, default='models/yolov5s.yaml',
|
|
|
|
|
help='model.yaml path')
|
|
|
|
|
parser.add_argument('--hyp', type=str, default='conf/hyp.scratch.yaml',
|
|
|
|
|
help='hyperparameters path')
|
2023-08-17 11:59:31 +08:00
|
|
|
parser.add_argument('--depth_factor', type=int, default=32,
|
|
|
|
|
help="Depth factor.")
|
|
|
|
|
parser.add_argument('--disable_cuda', action='store_true',
|
|
|
|
|
help="Disable CUDA.")
|
|
|
|
|
parser.add_argument('--gpu_id', type=int, default=0,
|
|
|
|
|
help="Select which gpu to use.")
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def get_parser_for_inference():
|
|
|
|
|
"""Return argument parser for inference."""
|
|
|
|
|
parser = argparse.ArgumentParser()
|
|
|
|
|
parser.add_argument('--mode', default='image', choices=['image', 'video'],
|
|
|
|
|
help="Inference image or video.")
|
|
|
|
|
parser.add_argument('--video',
|
|
|
|
|
help="Video path if you choose to inference video.")
|
|
|
|
|
parser.add_argument('--inference_slot', action='store_true', default=False,
|
|
|
|
|
help="Perform slot inference.")
|
|
|
|
|
parser.add_argument('--thresh', type=float, default=0.3,
|
|
|
|
|
help="Detection threshold.")
|
|
|
|
|
parser.add_argument('--save', action='store_true', default=True,
|
|
|
|
|
help="Save detection result to file.")
|
|
|
|
|
parser.add_argument('--yoloimg-size', type=int, default=640,
|
|
|
|
|
help='inference size (pixels)')
|
|
|
|
|
parser.add_argument('--dmprimg-size', type=int, default=640,
|
|
|
|
|
help='inference size (pixels)')
|
|
|
|
|
parser.add_argument('--augment', action='store_true',
|
|
|
|
|
help='augmented inference')
|
|
|
|
|
parser.add_argument('--conf-thres', type=float, default=0.25,
|
|
|
|
|
help='object confidence threshold')
|
|
|
|
|
parser.add_argument('--iou-thres', type=float, default=0.45,
|
|
|
|
|
help='IOU threshold for NMS')
|
2023-09-19 16:02:42 +08:00
|
|
|
parser.add_argument('--scale-ratio', type=float, default=0.5,
|
|
|
|
|
help='detected box scale ratio')
|
2023-08-17 11:59:31 +08:00
|
|
|
parser.add_argument('--ovlap-thres', type=float, default=0.6, help='overlap threshold for OBS')
|
|
|
|
|
parser.add_argument('--agnostic-nms', action='store_true',
|
|
|
|
|
help='class-agnostic NMS')
|
|
|
|
|
parser.add_argument('--classes', nargs='+', type=int,
|
|
|
|
|
help='filter by class: --class 0, or --class 0 2 3')
|
|
|
|
|
parser.add_argument('--dmpr-thresh', type=float, default=0.3,
|
|
|
|
|
help="Detection threshold.")
|
|
|
|
|
add_common_arguments(parser)
|
|
|
|
|
return parser
|
