diff --git a/abraia/detect.py b/abraia/detect.py index 4181dbe..e1681cb 100644 --- a/abraia/detect.py +++ b/abraia/detect.py @@ -169,10 +169,13 @@ def get_mask(row, box, img_width, img_height): return mask -def get_polygon(mask, origin): - """Calculates the bounding polygon based on the segmentation mask.""" - contours = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) - polygon = [[int(contour[0][0]), int(contour[0][1])] for contour in contours[0][0]] +def mask_to_polygon(mask, origin): + """Returns the largest bounding polygon based on the segmentation mask.""" + contours = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)[0] + polygon = [] + for contour in contours: + contour = contour.reshape(-1, 2) + polygon = list(contour) if contour.shape[0] > len(polygon) else polygon polygon = [(int(origin[0] + point[0]), int(origin[1] + point[1])) for point in polygon] return polygon @@ -211,9 +214,10 @@ def process_output(outputs, size, shape, classes, confidence, iou_threshold): for result in results: if len(outputs) == 2: + x1, y1, x2, y2 = result['box'] mask = result['mask'] @ output1 mask = get_mask(mask, (x1, y1, x2, y2), img_width, img_height) - result['polygon'] = get_polygon(mask, (x1, y1)) + result['polygon'] = mask_to_polygon(mask, (x1, y1)) result.pop('mask', None) return results @@ -277,7 +281,7 @@ def load_model(model_uri): if __name__ == '__main__': src = 'images/birds.jpg' - model_uri = 'https://api.abraia.me/files/multiple/camera/yolov8n.onnx' + model_uri = 'https://api.abraia.me/files/multiple/camera/yolov8n-seg.onnx' model = load_model(model_uri) @@ -290,13 +294,13 @@ def load_model(model_uri): img.show() - src = 'images/people-walking.mp4' + # src = 'images/people-walking.mp4' - def callback(img): - results = model.run(img) - objects = count_objects(results) - img = render_results(img, results) - return img + # def callback(img): + # results = model.run(img) + # objects = count_objects(results) + # img = render_results(img, results) + # return img - load_video(src, callback=callback) #, output='output.mp4') + # load_video(src, callback=callback) #, output='output.mp4') \ No newline at end of file