xiaoxue_admin
/
supervision
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Inference module now supports tracking detected objcets, extending process_roboflow_result to support that.

This commit is contained in:
ashishdatta 2023-10-07 10:24:26 -07:00
parent d30d80717c
commit cfb1e0ea60
1 changed files with 9 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

12
supervision/detection/utils.py
View File

@ -333,14 +333,15 @@ def extract_ultralytics_masks(yolov8_results) -> Optional[np.ndarray]:
def process_roboflow_result(
roboflow_result: dict,
) -> Tuple[np.ndarray, np.ndarray, np.ndarray, Optional[np.ndarray]]:
) -> Tuple[np.ndarray, np.ndarray, np.ndarray, Optional[np.ndarray], np.ndarray]:
if not roboflow_result["predictions"]:
return np.empty((0, 4)), np.empty(0), np.empty(0), None
return np.empty((0, 4)), np.empty(0), np.empty(0), None, None
xyxy = []
confidence = []
class_id = []
masks = []
tracker_ids = []
image_width = int(roboflow_result["image"]["width"])
image_height = int(roboflow_result["image"]["height"])
@ -359,6 +360,8 @@ def process_roboflow_result(
xyxy.append([x_min, y_min, x_max, y_max])
class_id.append(prediction["class_id"])
confidence.append(prediction["confidence"])
if "tracker_id" in prediction:
tracker_ids.append(prediction["tracker_id"])
elif len(prediction["points"]) >= 3:
polygon = np.array(
[[point["x"], point["y"]] for point in prediction["points"]], dtype=int
@ -368,13 +371,16 @@ def process_roboflow_result(
class_id.append(prediction["class_id"])
confidence.append(prediction["confidence"])
masks.append(mask)
if "tracker_id" in prediction:
tracker_ids.append(prediction["tracker_id"])
xyxy = np.array(xyxy) if len(xyxy) > 0 else np.empty((0, 4))
confidence = np.array(confidence) if len(confidence) > 0 else np.empty(0)
class_id = np.array(class_id).astype(int) if len(class_id) > 0 else np.empty(0)
masks = np.array(masks, dtype=bool) if len(masks) > 0 else None
tracker_id = np.array(tracker_ids).astype(int) if len(tracker_ids) > 0 else None
return xyxy, confidence, class_id, masks
return xyxy, confidence, class_id, masks, tracker_id
def move_boxes(xyxy: np.ndarray, offset: np.ndarray) -> np.ndarray: