import os
from PIL import Image
from openpyxl import load_workbook
from tqdm import tqdm
# Define the base path for the dataset
base_path = "/data5/Temp_ProjectData/QiJianProject/2022ZD0116304-202405/Images/CMB/OrginalData"

# Load the Excel workbook and select the active sheet
workbook_path = "/data1/Hatcher/model/project/MPKFNet/大规模数据库-图像.xlsx"  # Update this with the actual path to the Excel file
workbook = load_workbook(workbook_path)
sheet = workbook.active  # Assuming the active sheet is the one we need

# Function to get image dimensions
def get_image_dimensions(image_path):
    with Image.open(image_path) as img:
        return img.size  # Returns (width, height)

# Find the last row with data to continue from there
last_row = sheet.max_row + 1

# Function to update the Excel with new image data
def update_excel_with_images(base_path, start_row):
    # Iterate through all subfolders in the base directory
    a = 0
    for folder_name in tqdm(os.listdir(base_path)):
        folder_path = os.path.join(base_path, folder_name)
        if os.path.isdir(folder_path):  # Ensure it's a directory
            # Iterate through all images in the folder
   
            for image_name in os.listdir(folder_path):
                try:
                    if image_name.lower().endswith(('.png', '.jpg', '.jpeg', '.bmp', '.gif')):  # Check image extensions
                        image_path = os.path.join(folder_path, image_name)
                        dimensions = get_image_dimensions(image_path)
                        size_format = f"{dimensions[0]}x{dimensions[1]}"  # Format as 'WidthxHeight'

                        # Fill the Excel sheet
                        sheet.cell(row=start_row, column=1, value=start_row - 1)  # Sequence number
                        sheet.cell(row=start_row, column=2, value=image_path)     # Original image data
                        sheet.cell(row=start_row, column=3, value=folder_name)    # Label (folder name)
                        sheet.cell(row=start_row, column=4, value=size_format)    # Image size
                        sheet.cell(row=start_row, column=5, value=1)              # Data function type is always 1

                        # Increment the row index for the next entry
                        start_row += 1
                except:
                    pass

    return start_row
update_excel_with_images(base_path, last_row)
workbook.save(workbook_path)