Files
ds_task_recycling_project/backend/convert_csv_to_yolo.py
T
2025-07-17 00:03:03 +01:00

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)}")