urban_management/conf/config.py

"""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.")
    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')
    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