YOLOv5-Reward-driven-loss

YOLOv5 adaptation for screen estimation tasks with reward-driven loss function

<!-- Provide a quick summary of what the model is/does. -->

This model is a YOLOv5 model adapted for screen time estimation where object from each class can appear at most once per image. Loss function has been changed to follow a reward-driven mindset. Motivation and implementation decisions are explained in paper [will be linked as soon as it is published on IEEE].

Model Details

Model Description

<!-- Provide a longer summary of what this model is. -->

Other than usual YOLOv5 parameters, this implementation adds a few new ones:

    # 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')

Usage example

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()