import os
import shutil
import torch

# 设置根目录路径
root_directory = "/mnt/vos-s9gjtkm2/reid/dataset/imagenet"
val_directory = os.path.join(root_directory, "val")
val_images_directory = val_directory

# 读取验证集的 ground truth 文件
ground_truth_file = os.path.join(root_directory, "ILSVRC2012_devkit_t12", "data", "ILSVRC2012_validation_ground_truth.txt")
with open(ground_truth_file, "r") as f:
    val_labels = f.readlines()

# 读取 meta.bin 文件以获取类别信息
meta_file = os.path.join(root_directory, "meta.bin")
wnid_to_classes, val_wnids = torch.load(meta_file)

# 创建类别目录
for wnid in wnid_to_classes.keys():
    class_directory = os.path.join(val_directory, wnid)
    if not os.path.exists(class_directory):
        os.makedirs(class_directory)

# 将图片按类别分类
for i, label in enumerate(val_labels):
    label = label.strip()
    wnid = val_wnids[i]
    src_file = os.path.join(val_images_directory, f"ILSVRC2012_val_{i + 1:08d}.JPEG")
    dest_file = os.path.join(val_directory, wnid, f"ILSVRC2012_val_{i + 1:08d}.JPEG")
    shutil.move(src_file, dest_file)

print("Validation images have been classified by category.")