38 lines
1.4 KiB
Python
38 lines
1.4 KiB
Python
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)}") |