import os
import pandas as pd
from PIL import Image, ImageDraw, ImageFont
import random

# Define directories
base_dir = "dataset"
cars_dir = os.path.join(base_dir, "cars")
trees_dir = os.path.join(base_dir, "trees")
os.makedirs(cars_dir, exist_ok=True)
os.makedirs(trees_dir, exist_ok=True)

# Generate synthetic image with label
def generate_image(file_path, label, img_size=(64, 64)):
    bg_color = (random.randint(0, 255), random.randint(0, 255), random.randint(0, 255))
    image = Image.new("RGB", img_size, bg_color)
    draw = ImageDraw.Draw(image)
    try:
        font = ImageFont.truetype("arial.ttf", 20)
    except:
        font = ImageFont.load_default()
    text = label
    text_bbox = draw.textbbox((0, 0), text, font=font)
    text_width, text_height = text_bbox[2] - text_bbox[0], text_bbox[3] - text_bbox[1]
    text_position = ((img_size[0] - text_width) // 2, (img_size[1] - text_height) // 2)
    draw.text(text_position, text, fill="white", font=font)
    image.save(file_path)

# Generate images and collect data
num_images_per_class = 100
image_data = []
for i in range(num_images_per_class):
    file_name = f"car_{i}.jpg"
    file_path = os.path.join(cars_dir, file_name)
    generate_image(file_path, "Car")
    image_data.append({"image_path": file_path, "label": 0})
for i in range(num_images_per_class):
    file_name = f"tree_{i}.jpg"
    file_path = os.path.join(trees_dir, file_name)
    generate_image(file_path, "Tree")
    image_data.append({"image_path": file_path, "label": 1})

# Save to CSV
df = pd.DataFrame(image_data)
csv_path = "objects.csv"
df.to_csv(csv_path, index=False)
print(f"Created {csv_path} with {len(df)} entries.")

# Print sample
print("\nSample of objects.csv:")
print(df.head())