数据集处理

针对图像融合数据被包含在很多小文件夹中，将红外和RGB图片分别整理在两个文件夹中，这里针对TNO数据集

import os
import shutil

# 设置顶层文件夹路径
parent_dir = 'data/TNO_Image_Fusion_Dataset/'

# 创建保存红外和RGB图像的文件夹
ir_dir = 'data/TNO-after-processing/IR/'
rgb_dir = 'data/TNO-after-processing/RGB/'
os.makedirs(ir_dir, exist_ok=True)
os.makedirs(rgb_dir, exist_ok=True)


    
    
def process_dir(dir_path):
    #print("正在处理文件夹："+dir_path)
    for p in os.scandir(dir_path):
        if p.is_dir(): 
            has_img = False
            for f in os.scandir(p):
                if f.name.endswith('.bmp'):
                    has_img=True
                    break
            if has_img==True:
                read_img(p)
            process_dir(p.path) # 递归处理子文件夹
        elif p.is_file():
            return # 处理图像文件    
        
def read_img(subdir_path):
    if not os.path.isdir(subdir_path):
        return
    ir_img_lab = 0
    rgb_img_lab = 0
    # 在每个子文件夹内查找IR和VIS图像 
    for fname in os.listdir(subdir_path):
        fname = fname.split('.')[0]
        print("fname:"+fname)
        if fname.startswith(('IR' ,'ir','1','NIR','Ir')) or fname.endswith(('IR','ir','NIR','Ir')) :
            ir_img = os.path.join(subdir_path, fname+'.bmp')  
            ir_img_lab = 1        
        elif fname.startswith(('VIS' ,'vis','2','Vis')) or fname.endswith(("VIS",'vis','Vis')):
            rgb_img = os.path.join(subdir_path, fname+'.bmp')
            rgb_img_lab = 1
    if ir_img_lab==1 and rgb_img_lab==1:        
        # 构建新的匹配的文件名                
        new_fname = os.path.basename(subdir_path) + '_' + os.path.basename(ir_img)
        print("移动图片"+new_fname) 
        # 将图像复制到分类文件夹
        shutil.copy2(ir_img, os.path.join(ir_dir, new_fname))
        shutil.copy2(rgb_img, os.path.join(rgb_dir, new_fname))


if __name__ == '__main__': 
  
    root_dir = parent_dir 
    process_dir(root_dir)