YOLOv11 pose

Support YOLO
, ,
1

I am using the YOLOv11 pose model for cabbage head detection. After running the prediction with the best weights, the size of the keypoints is very large. How can I adjust the keypoint size to match the original?
This is the codes

#last codes 
import os
import torch
import cv2
import math
from ultralytics import YOLO

# Define the pixel-to-mm ratio (adjust using calibration)
pixel_to_mm_ratio = 1  # Example: 1 pixel = 1 mm (adjust accordingly)

# Function to calculate the Euclidean distance
def calculate_distance(point1, point2):
    return math.sqrt((point1[0] - point2[0])**2 + (point1[1] - point2[1])**2)

# Load the YOLO model
model = YOLO('best.pt')

# Define the output folder for predictions
output_folder = 'output'

# Ensure the output folder exists
os.makedirs(output_folder, exist_ok=True)

# Perform predictions
with torch.no_grad():
    results = model.predict(
        task="pose",
        source='image.jpg',
        max_det=150,
        conf=0.20,
        show_labels=False,
        show_conf=False,
        save=True,
        iou=0.6,
        agnostic_nms=True,
        device="0"
    )

# Process each prediction
for result in results:
    img_path = result.path  # Get image path
    img = cv2.imread(img_path)  # Load the image
    annotations = result.boxes.data.cpu().numpy()  # Get bounding boxes and pose points

    # Prepare a text file to save distances
    base_name = os.path.splitext(os.path.basename(img_path))[0]
    text_file_path = os.path.join(output_folder, f"{base_name}_distances.txt")

    with open(text_file_path, 'w') as f:
        for ann in annotations:
            # Bounding box coordinates
            x_min, y_min, x_max, y_max = ann[:4]
            center_x = (x_min + x_max) / 2
            center_y = (y_min + y_max) / 2
            center = (center_x, center_y)

            # Pose points (assuming these are stored from index 4 onwards in your annotation)
            try:
                points = ann[4:].reshape(-1, 2)  # Adjust if necessary for your model's output
            except ValueError:
                print(f"Pose points format incorrect for annotation: {ann}")
                continue

            # Step 1: Print the pose keypoints (pose points)
            f.write(f"Pose Keypoints: {points}\n")
            print(f"Pose Keypoints: {points}")

            # Step 2: Calculate distance between two specific points (center and the first pose point)
            if len(points) > 0:
                # Calculate distance from the center to the first pose point in pixels
                distance_px = calculate_distance(center, points[0])

                # Convert to millimeters (adjust the ratio as necessary)
                distance_mm = distance_px * pixel_to_mm_ratio

                # Write the calculated distance to the text file
                f.write(f"Distance from center to first pose point: {distance_mm:.2f} mm\n")
                print(f"Distance from center to first pose point: {distance_mm:.2f} mm")

                # Annotate the distance on the image
                cv2.putText(img, f"Dist: {distance_mm:.2f}mm", (int(center_x), int(center_y) - 10),
                            cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 1)

                # Optionally, draw a line from center to pose point
                cv2.line(img, (int(center_x), int(center_y)), (int(points[0][0]), int(points[0][1])), (0, 255, 0), 2)

    # Save the updated image with annotations
    save_path = os.path.join(output_folder, f"{base_name}_with_distances.jpg")
    cv2.imwrite(save_path, img)
    print(f"Processed image saved to {save_path}")
2

You can pass kpt_radius in model.predict(). Default is 5.

1 Like
3

Thank you so much for replying.