importosimportrandomimporttimefrompathlibimportPathimportshutilimporttkinterastkfromtkinterimportfiledialogfromloguruimportloggerimportxml.etree.ElementTreeasETclassAnalysisXML(object):'''清洗xml'''def__init__(self):root=tk.Tk()root.withdraw()root.attributes('-topmost',1)self.directory=filedialog.askdirectory()# 打开目录选择器root.destroy()logger.warning(f'路径选择:【{self.directory}】')defxml_img_split(self):'''分割图片和xml'''logger.info(f'---------------------------分割图片和xml------------------------')self.images_path=Path(self.directory).parent.joinpath('images')self.xml_labels_path=Path(self.directory).parent.joinpath('xml_labels')self.images_path.mkdir(parents=True,exist_ok=True)self.xml_labels_path.mkdir(parents=True,exist_ok=True)foriinPath(self.directory).iterdir():ifi.suffix=='.xml':new_path=self.xml_labels_path.joinpath(i.name)logger.debug(f'移动:【{i}】 -> 【{new_path}】')shutil.copy(str(i),str(new_path))ifi.suffixin('.jpg','.png'):new_path=self.images_path.joinpath(i.name)logger.debug(f'移动:【{i}】 -> 【{new_path}】')shutil.copy(str(i),str(new_path))defxml_to_txt(self):'''xml转txt'''logger.info(f'----------------------------正在将xml转为txt-----------------------')self.txt_labels=self.xml_labels_path.joinpath('labels')# 替换为实际的输出TXT文件夹路径os.makedirs(self.txt_labels,exist_ok=True)names_set=set()forfilenameinos.listdir(self.xml_labels_path):iffilename.endswith('.xml'):tree=ET.parse(os.path.join(self.xml_labels_path,filename))root=tree.getroot()forobjinroot.findall('object'):name=obj.find('name').text names_set.add(name)# 输出所有的namecategories=[]fornameinnames_set:categories.append(name)logger.success(f'标注的内容names:【{categories}】')category_to_index={category:indexforindex,categoryinenumerate(categories)}# 遍历输入文件夹中的所有XML文件forfilenameinos.listdir(self.xml_labels_path):iffilename.endswith('.xml'):xml_path=os.path.join(self.xml_labels_path,filename)logger.warning(f'正在处理:【{xml_path}】')# 解析XML文件tree=ET.parse(xml_path)root=tree.getroot()# 提取图像的尺寸size=root.find('size')width=int(size.find('width').text)height=int(size.find('height').text)# 存储name和对应的归一化坐标objects=[]# 遍历XML中的object标签forobjinroot.findall('object'):name=obj.find('name').textifnameincategory_to_index:category_index=category_to_index[name]else:continue# 如果name不在指定类别中，跳过该objectbndbox=obj.find('bndbox')xmin=int(bndbox.find('xmin').text)ymin=int(bndbox.find('ymin').text)xmax=int(bndbox.find('xmax').text)ymax=int(bndbox.find('ymax').text)# 转换为中心点坐标和宽高x_center=(xmin+xmax)/2.0y_center=(ymin+ymax)/2.0w=xmax-xmin h=ymax-ymin# 归一化x=x_center/width y=y_center/height w=w/width h=h/height objects.append(f"{category_index}{x:.6f}{y:.6f}{w:.6f}{h:.6f}")# 输出结果到对应的TXT文件txt_filename=os.path.splitext(filename)[0]+'.txt'txt_path=os.path.join(self.txt_labels,txt_filename)withopen(txt_path,'w')asf:forobjinobjects:f.write(obj+'\n')defto_dataset(self,test_ratio):'''整理为dataset'''output_folder=os.path.join(os.path.dirname(self.directory),'datasets')input_image_folder=self.images_path input_label_folder=self.txt_labels train_images_folder=os.path.join(output_folder,'train','images')train_labels_folder=os.path.join(output_folder,'train','labels')val_images_folder=os.path.join(output_folder,'val','images')val_labels_folder=os.path.join(output_folder,'val','labels')os.makedirs(train_images_folder,exist_ok=True)os.makedirs(train_labels_folder,exist_ok=True)os.makedirs(val_images_folder,exist_ok=True)os.makedirs(val_labels_folder,exist_ok=True)# 获取所有图像文件列表images=[fforfinos.listdir(input_image_folder)iff.endswith('.jpg')orf.endswith('.png')]# 随机打乱图像文件列表random.shuffle(images)# 计算验证集的数量val_size=int(len(images)*test_ratio)# 划分验证集和训练集val_images=images[:val_size]train_images=images[val_size:]# 复制验证集图像和标签forimageinval_images:label=os.path.splitext(image)[0]+'.txt'ifos.path.exists(os.path.join(input_label_folder,label)):shutil.copy(os.path.join(input_image_folder,image),os.path.join(val_images_folder,image))shutil.copy(os.path.join(input_label_folder,label),os.path.join(val_labels_folder,label))logger.debug(f'【{os.path.join(input_image_folder,image)}】 --> 【{os.path.join(val_images_folder,image)}】')logger.success(f'【{os.path.join(input_label_folder,label)}】 --> 【{os.path.join(val_labels_folder,label)}】')else:logger.error(f"Warning: Label file{label}not found for image{image}")# 复制训练集图像和标签forimageintrain_images:label=os.path.splitext(image)[0]+'.txt'ifos.path.exists(os.path.join(input_label_folder,label)):shutil.copy(os.path.join(input_image_folder,image),os.path.join(train_images_folder,image))shutil.copy(os.path.join(input_label_folder,label),os.path.join(train_labels_folder,label))logger.debug(f'【{os.path.join(input_image_folder,image)}】 --> 【{os.path.join(train_images_folder,image)}】')logger.success(f'【{os.path.join(input_label_folder,label)}】 --> 【{os.path.join(train_labels_folder,label)}】')else:logger.error(f"Warning: Label file{label}not found for image{image}")defstart(self):'''启动'''time.sleep(1)self.xml_img_split()time.sleep(1)self.xml_to_txt()time.sleep(1)self.to_dataset(0.2)if__name__=='__main__':base_dir=os.path.dirname(__file__)log_path=os.path.join(base_dir,'log.log')ifos.path.exists(log_path):os.unlink(log_path)logger.add(log_path)print('...第一层文件夹')print(' -->第二层文件夹↓')print(' -->[xml和img混合文件夹]')print('\n')status=input('请确认xml和图片在同一个文件夹(99:确认)(任意值:取消):')ifstatusin(99,'99'):a=AnalysisXML()a.start()logger.success('系统完成')foriin(3,2,1):time.sleep(1)logger.success(f'{i}/秒')else:logger.error('系统退出!')foriin(3,2,1):time.sleep(1)logger.error(f'{i}/秒')
