Search code examples
pythonobject-detectionyolov5

YOLO V5 custom object detection : csv output

I managed to train the YOLO V5 model on my customed data and I'm having great results. I was just wondering how I could export the bonding boxes in a csv or txt file in which I'd have the coordinates and the score of prediction. I've read that I'll need to change the detect.py file but I just don't know how.

Solution

  • in the line 159 at https://github.com/ultralytics/yolov5/blob/master/detect.py

    there is an argument to save txt files, usage shoud be something like:

    python detect.py --weights yolov5s.pt --save-txt det.txt

    the other option is to:

    import pandas as pd
detections = []

...
detections.append(line)

... 
#condition to save
df = pd.DataFrame.from_dict( {'detections': detections})
df.to_csv("detections.csv")