annotation_format_checker.py

import cv2
import matplotlib.pyplot as plt
import os

def check_yolo_annotations(image_dir, label_dir):
    """
    Checks if the conversion from VOC to YOLO format is correct by plotting bounding boxes on the images.

    Parameters:
    image_dir (str): Directory path to the images.
    label_dir (str): Directory path to the YOLO format annotations.
    """
    # Get the first image file
    image_files = sorted(os.listdir(image_dir))
    first_image_file = image_files[0]

    # Construct paths for the image and its corresponding label
    image_path = os.path.join(image_dir, first_image_file)
    label_path = os.path.join(label_dir, os.path.splitext(first_image_file)[0] + '.txt')

    # Load the image using OpenCV
    image = cv2.imread(image_path)
    # Convert the image from BGR (OpenCV default) to RGB (matplotlib default)
    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

    # Read the label file to get bounding box information
    with open(label_path, 'r') as f:
        lines = f.readlines()

    # Plot the bounding box on the image
    for line in lines:
        # Parse the label file line to extract bounding box information
        class_id, x_center, y_center, width, height = map(float, line.strip().split())
        img_height, img_width, _ = image.shape

        # Convert YOLO format to bounding box format
        x_center *= img_width
        y_center *= img_height
        width *= img_width
        height *= img_height

        # Calculate the top-left and bottom-right coordinates of the bounding box
        x1 = int(x_center - width / 2)
        y1 = int(y_center - height / 2)
        x2 = int(x_center + width / 2)
        y2 = int(y_center + height / 2)

        # Draw the bounding box on the image using a green rectangle
        cv2.rectangle(image, (x1, y1), (x2, y2), (0, 255, 0), 2)

    # Display the image with bounding box using matplotlib
    plt.imshow(image)
    plt.axis('off')  # Hide the axis
    plt.show()  # Display the image

if __name__ == "__main__":
    check_yolo_annotations('Dataset/images', 'Dataset/yolo_annotations')