File size: 432 Bytes
39a7537 |
1 2 3 4 5 6 7 8 9 10 11 |
CLASSES = {0:"body", 1:"face", 2:"frame", 3:"text"} # confirmed by user
def get_yolo_predictions(yolo_model,img_path:str):
results = yolo_model.predict(source=img_path, device='cuda')
dets = {"body":[], "face":[], "frame":[], "text":[]}
for box in results[0].boxes:
c = int(box.cls[0])
x1, y1, x2, y2 = map(int, box.xyxy[0])
dets[CLASSES[c]].append([x1,y1,x2,y2])
return dets |