import os
import pyheif
from PIL import Image

def convert_heic_to_jpeg_and_remove(source_folder):
    converted_count = 0
    for filename in os.listdir(source_folder):
        if filename.lower().endswith(".heic"):
            heic_path = os.path.join(source_folder, filename)
            jpeg_filename = os.path.splitext(filename)[0] + ".jpeg"
            jpeg_path = os.path.join(source_folder, jpeg_filename)

            try:
                heif_file = pyheif.read(heic_path)
                image = Image.frombytes(
                    heif_file.mode,
                    heif_file.size,
                    heif_file.data,
                    "raw",
                    heif_file.mode,
                    heif_file.stride,
                )
                image.save(jpeg_path, "JPEG")
                os.remove(heic_path)
                converted_count += 1
            except ValueError as e:
                print(f"Skipping {heic_path}: {e}")

    total_files = len(os.listdir(source_folder))
    return converted_count, total_files

def convert_png_to_jpeg(png_dir, jpeg_dir):
    if not os.path.exists(jpeg_dir):
        os.makedirs(jpeg_dir)

    for filename in os.listdir(png_dir):
        if filename.endswith('.png'):
            png_file = os.path.join(png_dir, filename)
            img = Image.open(png_file)
            rgb_img = img.convert('RGB')
            jpeg_file = os.path.join(jpeg_dir, filename.replace('.png', '.jpg'))
            rgb_img.save(jpeg_file, 'JPEG')

if __name__ == "__main__":
    source_folder = "datasets/imagedata"
    converted_count, total_files = convert_heic_to_jpeg_and_remove(source_folder)
    print(f"Total HEIC files converted: {converted_count}")
    print(f"Total number of files in the folder: {total_files}")

    png_directory = 'datasets/imagedata'
    jpeg_directory = 'datasets/imagedata/images'
    convert_png_to_jpeg(png_directory, jpeg_directory)