# !/usr/bin/env python
# -*- coding:utf-8 -*-
# @Time : 2024.03
# @Author : 绿色羽毛
# @Email : lvseyumao@foxmail.com
# @Blog : https://blog.csdn.net/ViatorSun
# @Note :
import os
import cv2
import os.path as osp
from resnet_classify import *
from utils import *
# 替换文件后缀，并输出文件名
def change_suffle(file_path, new_suffle=".txt"):
base_name, file_extension = os.path.splitext(file_path)
new_file = base_name + new_suffle # 对应标注文明名
file_name = osp.basename(base_name)
return file_name, new_file
# 读取yolo标注信息，按类别 将其转化为像素坐标，返回字典格式
def read_labels(label_path, image_shape):
labels_pixel = [] # 像素格式
labels_yolo = [] # 原标注格式
with open(label_path,'r') as f:
for line in f.readlines():
line_split = line.strip().split(' ')
if len(line_split)!=5:
continue
cls_id=int(line_split[0])
box=list(map(float,line[1:].strip().split(' ')))
half_box_w=box[2]*0.5
half_box_h=box[3]*0.5
xmin=int((box[0]-half_box_w)*image_shape[1])
xmax=int((box[0]+half_box_w)*image_shape[1])
ymin=int((box[1]-half_box_h)*image_shape[0])
ymax=int((box[1]+half_box_h)*image_shape[0])
labels_pixel.append([cls_id, xmin, ymin, xmax, ymax])
labels_yolo.append(line)
return labels_pixel, labels_yolo
def find_key_by_value(dictionary, value):
for key, val in dictionary.items():
if val == value:
return key
# 如果未找到匹配的键，则返回 None
return None
class change_class(object):
def __init__(self, configs:dict):
self.model_lst = {}
self.old_labels = {}
self.new_labels = {}
for conf in configs.keys():
self.model_lst[conf] = ResnetClassify(configs[conf][0])
self.old_labels[conf] = configs[conf][1]
self.new_labels[conf] = configs[conf][2]
def change(self, data_path, output_path):
# 创建输出文件夹
folder = osp.basename(data_path)
save_path = osp.join(output_path, folder)
if not osp.exists(save_path):
os.makedirs(save_path)
# 遍历 data_path 下所有的图片
file_lst = scan_files(data_path)
for img_path in file_lst:
file_name, txt_path = change_suffle(img_path)
img = cv2.imread(img_path)
image_shape = img.shape
labels_pixel, labels_yolo = read_labels(txt_path, image_shape[:-1]) # 标注信息转换
new_labels = []
for pixel_idx, yolo_idx in zip(labels_pixel, labels_yolo):
if pixel_idx[0] in self.old_labels.values():
model_name = find_key_by_value(self.old_labels, pixel_idx[0])
xmin, ymin, xmax, ymax = pixel_idx[1:]
img_roi = img[ymin:ymax, xmin:xmax]
new_cls = self.model_lst[model_name].predict(img_roi) + self.new_labels[model_name]
labels = str(new_cls) + yolo_idx[len(str(pixel_idx[0])):]
new_labels.append(labels)
else:
new_labels.append(yolo_idx)
# 拷贝数据
new_img_path = osp.join(save_path, osp.basename(img_path))
new_txt_path = osp.join(save_path, file_name + ".txt")
os.system(f"cp {img_path} {new_img_path}")
# 打开文件以写入模式（'w'）
with open(new_txt_path, 'w') as f:
# 将数据逐行写入文件
for line in new_labels:
f.write(line)
print(f">> handle {file_name}")
if __name__ == '__main__':
# 模型权重 旧标签 替换起始标签
resnet_dict = {"yb": ["/media/sun/DataYZ/Project/GenerateData/change_bbox_clas/weights/压板/yb_cls.torchscript", 6, 6],
"kk": ["/media/sun/DataYZ/Project/GenerateData/change_bbox_clas/weights/空开/kk_cls.torchscript", 7, 9],
"zsd": ["/media/sun/DataYZ/Project/GenerateData/change_bbox_clas/weights/指示灯/zsd_cls.torchscript", 8, 11],
"xn": ["/media/sun/DataYZ/Project/GenerateData/change_bbox_clas/weights/旋钮/xn_cls.torchscript", 9, 13]}
change_model = change_class(resnet_dict)
data_path = "/media/sun/DataYZ/DataSet/2024年03月13日开始_压板空开等改为带状态/第一波"
save_path = "./output"
change_model.change(data_path, save_path)
