convert_annotations.py

import os
import xml.etree.ElementTree as ET

# Define paths
annotations_path = "archive/annotations"
output_path = "archive/labels"
images_path = "archive/images"

# Class mapping (adjust according to your dataset)
class_mapping = {
    "with_mask": 0,
    "without_mask": 1,
    "mask_weared_incorrect": 2
}

# Create labels directory if it doesn't exist
os.makedirs(output_path, exist_ok=True)

# Function to convert XML to YOLO format
def convert_to_yolo(xml_file, image_width, image_height):
    tree = ET.parse(xml_file)
    root = tree.getroot()
    yolo_annotations = []
    
    for obj in root.findall("object"):
        class_name = obj.find("name").text
        class_id = class_mapping[class_name]
        
        # Get bounding box coordinates
        bndbox = obj.find("bndbox")
        xmin = int(bndbox.find("xmin").text)
        ymin = int(bndbox.find("ymin").text)
        xmax = int(bndbox.find("xmax").text)
        ymax = int(bndbox.find("ymax").text)
        
        # Convert to YOLO format
        x_center = (xmin + xmax) / 2.0 / image_width
        y_center = (ymin + ymax) / 2.0 / image_height
        width = (xmax - xmin) / image_width
        height = (ymax - ymin) / image_height
        yolo_annotations.append(f"{class_id} {x_center:.6f} {y_center:.6f} {width:.6f} {height:.6f}")
    
    return yolo_annotations

# Process each annotation file
for annotation_file in os.listdir(annotations_path):
    if annotation_file.endswith(".xml"):
        xml_path = os.path.join(annotations_path, annotation_file)
        image_name = annotation_file.replace(".xml", ".png")
        image_path = os.path.join(images_path, image_name)
        
        # Check if the corresponding image exists
        if os.path.exists(image_path):
            # Get image dimensions
            import cv2
            img = cv2.imread(image_path)
            image_height, image_width, _ = img.shape
            
            # Convert annotations
            yolo_data = convert_to_yolo(xml_path, image_width, image_height)
            
            # Write to YOLO label file
            label_file = os.path.join(output_path, annotation_file.replace(".xml", ".txt"))
            with open(label_file, "w") as f:
                f.write("\n".join(yolo_data))