Yolov8: finding the corners of a polygon segmentation

Question

I’m trying to find the corners of a polygon segmentation that was made with Yolov8, as in this image: chessboard segmentation

This is my code:

model_trained = YOLO("runs/segment/yolov8n-seg_chessboard/weights/best.pt")

results = model_trained.predict(source="1.jpgresized.jpg", line_thickness=2, save_txt=True, save=True)

masks = results[0].masks # Masks object

masks.segments[0] # a numpy array of

I'm not able to figure out how to get the four corners of the segmentation out of this array.

Thanks!

Answer 1

Notice, that this could involve quite a lot of fine-tuning for you particular case. The idea here is to pass the segmentation mask to goodFeaturesToTrack which finds strong corners in it. Then you pick the 4 best candidates. Finally they are plotted on the original image.

from ultralytics.yolo.engine.model import YOLO
import cv2


def on_predict_batch_end(predictor):
    # results -> List[batch_size]
    path, im, im0s, vid_cap, s = predictor.batch
    predictor.results = zip(predictor.results, im0s)


model = YOLO("yolov8n-seg.pt")
model.add_callback("on_predict_batch_end", on_predict_batch_end)
results = model.predict(source="0", show=True, stream=True, classes=67)
for i, (result, im0) in enumerate(results):
    Masks = result.masks
    if Masks is not None:
        for mask in Masks.masks:
            x = mask.cpu().numpy()
            corners = cv2.goodFeaturesToTrack(x, 4, 0.5 , 50)  # get 4 corners
            for corner in corners:  # plot the corners on the original image
                x,y = corner.ravel()
                cv2.circle(im0,(int(x), int(y)),5,(0, 0, 255),-1)