import os
import shutil

# 数据集路径
dataset_path = '/mnt/Study/DataSet/DataSet/fgvc-aircraft-2013b/fgvc-aircraft-2013b/data/images'
output_path = '/mnt/Study/DataSet/DataSet/fgvc-aircraft-2013b/fgvc-aircraft-2013b/data/sorted_images'

# 类别文件，例如 'images_variant_trainval.txt'
labels_file = '/mnt/Study/DataSet/DataSet/fgvc-aircraft-2013b/fgvc-aircraft-2013b/data/images_variant_test.txt'

# 创建输出文件夹
if not os.path.exists(output_path):
    os.makedirs(output_path)

# 读取类别文件
with open(labels_file, 'r') as f:
    lines = f.readlines()

count = 0

for line in lines:
    count += 1
    print(f'Processing image {count}/{len(lines)}', end='\r')
    parts = line.strip().split(' ')
    image_name = parts[0] + '.jpg'
    category = '_'.join(parts[1:]).replace('/', '_')

    # 创建类别文件夹
    category_path = os.path.join(output_path, category)
    if not os.path.exists(category_path):
        os.makedirs(category_path)

    # 移动图像到对应类别文件夹
    src = os.path.join(dataset_path, image_name)
    dst = os.path.join(category_path, image_name)
    if os.path.exists(src):
        shutil.move(src, dst)
    else:
        print(f'Image {image_name} not found!')

print("Images have been sorted into folders by category.")
