import os
from PIL import Image
from openpyxl import load_workbook
from tqdm import tqdm

# Define the base paths for images and labels
base_path = "/data5/Temp_ProjectData/QiJianProject/2022ZD0116304-202405/Images/CDB"
image_folder = os.path.join(base_path, "JPEGImages")
label_folder = os.path.join(base_path, "txt")

# Load the Excel workbook and select the first 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 first sheet is the target one

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

# Start filling the Excel from row 2 (row 1 is the header)
row_index = 2

# Iterate over images and corresponding label files
for image_name in tqdm(os.listdir(image_folder)):
    if image_name.endswith(".jpg"):  # Ensure the file is a JPEG image
        # Construct full paths for the image and label
        image_path = os.path.join(image_folder, image_name)
        label_name = os.path.splitext(image_name)[0] + ".txt"
        label_path = os.path.join(label_folder, label_name)

        # Get image dimensions
        dimensions = get_image_dimensions(image_path)

        # Fill in the Excel sheet
        sheet.cell(row=row_index, column=1, value=row_index - 1)  # Sequence number
        sheet.cell(row=row_index, column=2, value=image_path)     # Original image data
        sheet.cell(row=row_index, column=3, value=label_path)     # Label
        sheet.cell(row=row_index, column=4, value=f"{dimensions[0]}x{dimensions[1]}")  # Image size
        sheet.cell(row=row_index, column=5, value=0)              # Data function type

        # Increment the row index for the next entry
        row_index += 1
      

# Save the updated workbook
workbook.save(workbook_path)