# Arguments added for screen time estimation task
    parser.add_argument('--iou_treshold', type=float, default=0.0, help='IoU treshold for deciding if prediction is good')
    parser.add_argument('--confidence_treshold', type=float, default=0.0, help='Confidence treshold for deciding if prediction is good')
    parser.add_argument('--loss_function', type=str, choices=['ordinary', 'ignore', 'motivate'], default='motivate', help='Which loss function to use')
    parser.add_argument('--class_wise_nms', type=bool, default=True, help='To use class wise non max surpression or not')

from train import run


def main():

    conf_th = 0.05
    iou_th = 0.80
    data = 'sign_language.yaml'
    loss_function = 'ordinary'

    run(imgsz=640, epochs=100, data=data, weights='', cfg='models/yolov5s.yaml', patience=0, confidence_treshold=conf_th, iou_treshold=iou_th, loss_function=loss_function, class_wise_nms=True)

if __name__ == "__main__":
    main()