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