I trained my custom yolo dataset, now my requirement is to feed an image to the model and get back a json with founded classes with scores and bounding boxes.
When I run detect.py, it looks that there is no option to configure this script to my needs, and I can only get back an image with bounding box and score placed on it.
How did you solve this problem? Did you took parts of detect.py and modified to your needs, or is there probably a solution to it?
I have found the answer to it.
Once you have your model, you need this script to make it work and send the answer to a calling program. I had just some problems with the paths, but in the end figured it out. If you think it can be better, please let me know:
import os
import torch
# Important: Path here is taken not from the folder where the script is positioned, but from the root project folder
model = torch.hub.load(os.getcwd()+'/yolov5','custom', path='./yolov5/runs/train/lemon3/weights/best.pt', source='local', force_reload=True)
results = model("./dataset/test/images/fridge2.jpg")
# cord_thres contains 4 corners of bounding box, 5th array parameter is confidence score
labels, cord_thres = results.xyxyn[0][:, -1].numpy(), results.xyxyn[0][:, :-1].numpy()
print(labels, cord_thres)