import os
import csv
import matplotlib.pyplot as plt
import matplotlib.image as mpimg

# Set the path to your extracted directory
images_path = './numbers'
output_csv = 'labels2.csv'

# Create/open the CSV file for saving labels
with open(output_csv, mode='w', newline='') as file:
    writer = csv.writer(file)
    writer.writerow(['Image Name', 'Label'])  # Header row

    for image_name in os.listdir(images_path):
        image_path = os.path.join(images_path, image_name)
        try:
            img = mpimg.imread(image_path)
            plt.imshow(img)
            plt.axis('off')
            plt.show()

            label = input(f'Enter the label for {image_name}: ')

            writer.writerow([image_name, label])
            print(f'Label saved for {image_name}: {label}')

        except Exception as e:
            print(f"Error processing {image_name}: {e}")