Autolabeling upload from SDK doesn't preserve label positions

[tomkr000]

I'm using a pretrained model to help with autolabeling keypoints. It works decently well, however when I use the SDK to upload, the points end up in the wrong places.

Here is the code:

from trainyolo.client import Client, Project
from trainyolo.utils.yolov8 import format_boxes, format_masks, format_keypoints
from ultralytics import YOLO
from urllib.request import urlretrieve
import os
from tqdm import tqdm

# initialize client
API_KEY = "xxxxxxxxxxxxxxx"
PROJECT_NAME = "mouse_cropped_2"
client = Client(API_KEY)

project = Project.get_by_name(client, PROJECT_NAME)
samples = project.get_samples(filter='UNLABELED')

# initialize YOLOv8 model
MODEL_PATH = "/home/ubuntu/code/notebooks/rodent_behavior/12e19d97-a248-4fa9-9627-116e0877db89.pt"
model = YOLO(MODEL_PATH)

# create a temporary image location
image_loc = './images'
os.makedirs(image_loc, exist_ok=True)

for sample in tqdm(samples):

  # download image
  image_path = os.path.join(image_loc, sample.name)
  if not os.path.exists(image_path):
      urlretrieve(sample.asset['url'], image_path)

  # forward image

  # img = cv2.imread(image_path)
  # img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
  detections = model.predict(source=image_path)

  # # OPTION 1: FOR DETECTION MODEL
  # boxes, cls = detections[0].boxes.xyxy.cpu().numpy(), detections[0].boxes.cls.cpu().numpy()
  # sample.prediction = format_boxes(boxes, cls)

  # # OPTION 2: FOR SEGMENTATION MODEL
  # masks, cls = detections[0].masks.cpu().numpy(), detections[0].boxes.cls.cpu().numpy()
  # sample.prediction = format_masks(masks, cls)

  boxes = detections[0].boxes.xyxy.cpu().numpy()
  keypoints = detections[0].keypoints.xy.cpu().numpy()
  cls = detections[0].boxes.cls.cpu().numpy()

  sample.prediction = format_keypoints(boxes, keypoints, cls)

The keypoints from the last image run through inference (I didn't run inference a second time, this is directly what was uploaded with sample.prediction):



And what they look like in the labeling window:

I also tried a couple different ways of loading the image, ex: through cv2 (commented out) but no luck.

Any thoughts on what is happening / what I'm doing wrong?

[davyneven]

Dear tomkr000, there is indeed something wrong here. I'll troubleshoot this asap. In the meantime, you can use the 'Auto Label' function (magic wand), which generates predictions using your uploaded model.

[davyneven]

Probably the issue is the way you extract the keypoints, it should be like this:

keypoints = detections[0].keypoints.data.cpu().numpy()

That should normally fix your issues.

[tomkr000]

That fixed it thank you!