import os
import shutil
import json
from PIL import Image

image_folder = "/mnt/petrelfs/zhuchenglin/LLaVA/playground/data/LLaVA-Pretrain/images"
target_folder = "/mnt/petrelfs/zhuchenglin/llava_gen_200k"
annotations_path = "/mnt/petrelfs/zhuchenglin/LLaVA/playground/data/LLaVA-Pretrain/llava_gen_200k.json"
with open(annotations_path, "r") as f:
    annotations = json.load(f)

for index, annotation in enumerate(annotations):
    print(index)
    folder_index_source1 = 700 + (index // 10000)
    folder_index_source2 = 680 + (index // 10000)
    folder_index_target = 700 + (index // 10000)
    subfolder_source1 = f"{folder_index_source1:05d}"
    subfolder_source2 = f"{folder_index_source2:05d}"
    subfolder_target = f"{folder_index_target:05d}"
    source1_image_name = f"{folder_index_source1:05d}{index % 10000:04d}.jpg"
    source2_image_name = f"{folder_index_source2:05d}{index % 10000:04d}.jpg"
    target_image_name = f"{folder_index_target:05d}{index % 10000:04d}.jpg"
    source1_image_path = os.path.join(
        image_folder, subfolder_source1, source1_image_name
    )
    source2_image_path = os.path.join(
        image_folder, subfolder_source2, source2_image_name
    )
    target_image_path = os.path.join(target_folder, subfolder_target, target_image_name)
    if not os.path.exists(os.path.join(target_folder, subfolder_target)):
        os.makedirs(os.path.join(target_folder, subfolder_target))
    with Image.open(source1_image_path) as source1_img:
        resized_img = source1_img.resize((336, 336))
        resized_img.save(target_image_path)