import pandas as pd
from pathlib import Path
import os


def csv_to_yolo(csv_path, output_dir):
    # Create output directory if it doesn't exist
    Path(output_dir).mkdir(parents=True, exist_ok=True)
    
    df = pd.read_csv(csv_path)
    
    for filename in df['filename'].unique():
        img_data = df[df['filename'] == filename].iloc[0]
        img_w, img_h = img_data['img_width'], img_data['img_height']
        
        yolo_lines = []
        for _, row in df[df['filename'] == filename].iterrows():
            # Convert absolute to normalized coordinates
            x_center = ((row['x1'] + row['x2']) / 2) / img_w
            y_center = ((row['y1'] + row['y2']) / 2) / img_h
            width = abs(row['x2'] - row['x1']) / img_w
            height = abs(row['y2'] - row['y1']) / img_h
            
            yolo_lines.append(f"0 {x_center:.6f} {y_center:.6f} {width:.6f} {height:.6f}")
        
        # Save as YOLO .txt file
        txt_path = Path(output_dir) / f"{Path(filename).stem}.txt"
        with open(txt_path, 'w') as f:
            f.write("\n".join(yolo_lines))
    print(f"Successfully converted CSV to YOLO format in {output_dir}")

# Error handling
try:
    csv_to_yolo("annotations.csv", "yolo_labels")
except FileNotFoundError:
    print("Error: annotations.csv not found. Please check the file path.")
except Exception as e:
    print(f"An error occurred: {str(e)}")