yolo plotting utility - Githubissues

import os
from PIL import Image
import matplotlib.pyplot as plt
import matplotlib.patches as patches

def plot_yolo_bounding_boxes(path_to_png, path_to_txt, category_map=None):
    # Open the image using PIL
    image = Image.open(path_to_png)
    image_width, image_height = image.size

    # Define colors for different category IDs
    category_colors = {1: 'red', 2: 'green', 3: 'blue', 4: 'yellow'}  # Update with your own category IDs and colors

    # Load the bounding box coordinates from the .txt file
    with open(path_to_txt, "r") as txt_file:
        for line in txt_file:
            label_parts = line.strip().split(" ")
            category_id = int(label_parts[0])
            x_center, y_center, bbox_width, bbox_height = map(float, label_parts[1:])
            x_min = max(0, int((x_center - bbox_width / 2) * image_width))
            y_min = max(0, int((y_center - bbox_height / 2) * image_height))
            x_max = min(image_width, int((x_center + bbox_width / 2) * image_width))
            y_max = min(image_height, int((y_center + bbox_height / 2) * image_height))

            # Get color for current category ID
            bbox_color = category_colors.get(category_id, 'red')  # Default to red if category ID not found

            # Plot the bounding box on the image with category-specific color
            rect = patches.Rectangle((x_min, y_min), (x_max - x_min), (y_max - y_min),
                                      fill=None, edgecolor=bbox_color, linewidth=1)
            plt.gca().add_patch(rect)
            if category_map is not None:
                category_name = category_map[category_id]
            else:
                category_name = f"Class {category_id}"
            plt.text(x_min, y_min, category_name, color=bbox_color, fontsize=6, weight='bold')

    # Show the image with bounding boxes
    plt.imshow(image)
    plt.axis('off')
    plt.show()

# Example usage:
path_to_png = "image1.jpeg"     # Path to the PNG image
path_to_txt = "image1.txt"     # Path to the YOLO formatted .txt file
plot_yolo_bounding_boxes(path_to_png, path_to_txt)

# Or
# Example usage:
fig, axs = plt.subplots(4, 4, figsize=(15, 15))  # Create a 4x4 grid of subplots
axs = axs.ravel()  # Flatten the subplots array

for i in range(16):
    path_to_png = f"example_{i+1}.png"     # Path to the PNG image
    path_to_txt = f"example_{i+1}.txt"     # Path to the YOLO formatted .txt file
    axs[i].set_title(f"Example {i+1}")
    plot_yolo_bounding_boxes(path_to_png, path_to_txt)

plt.tight_layout()
plt.show()
benhuff / random_geospatial

yolo plotting utility #5
