只转单个类别
import os
import cv2
# 传入图像地址,统计某张照片中的所有像素值
def countpixel(rootdir):
assert os.path.exists(rootdir), f'{rootdir} not exist.'
arr = [] # 查看了tunnelmask的像素值
for file in os.listdir(rootdir):
img = cv2.imread(os.path.join(rootdir, file))
height, width = img.shape[:-1]
for h in range(height):
for w in range(width):
pixel = img[h, w].tolist()
if pixel not in arr:
arr.append(pixel)
print(f'arr={arr}')
return arr
# 从一张图中找到[1, 1, 1]出现的x_min x_max y_min y_max
def position_lwd(imagepath, cls_list):
image = cv2.imread(imagepath)
height, width = image.shape[:-1]
x = [[]]
y = [[]]
cnt = 0 # 实例个数
for h in range(height):
for w in range(width):
pixel = image[h, w].tolist()
if pixel != [0, 0, 0]: # 如果不为背景
if w not in x:
x.append(w)
if h not in y:
y.append(h)
# 找出x_min x_max y_min y_max
x_min = min(x)
x_max = max(x)
y_min = min(y)
y_max = max(y)
return (x_min, x_max, y_min, y_max)
# 归一化处理
def convert(size, box):
# 该函数将xmin,ymin,xmax,ymax转为x,y,w,h中心点坐标和宽高
dw = 1. / (size[0])
dh = 1. / (size[1])
x = (box[0] + box[1]) / 2.0
y = (box[2] + box[3]) / 2.0
w = box[1] - box[0]
h = box[3] - box[2]
x = x * dw
w = w * dw
y = y * dh
h = h * dh
return (x, y, w, h)
def convert_annotation(rootdir, out_dir, cls_list):
assert os.path.exists(rootdir), f'{rootdir} not exist.'
if not os.path.exists(out_dir):
os.makedirs(out_dir)
filelist = os.listdir(rootdir)
for file in filelist:
filename = os.path.splitext(file)[0]
file_path = os.path.join(rootdir, file)
file_cv = cv2.imread(file_path)
h, w = file_cv.shape[:-1]
# 输出标签的地址
out_path = os.path.join(out_dir, filename + '.txt')
out_file = open(out_path, 'w')
# 开始转换
box = position_lwd(file_path, cls_list) # 多类别转换
bb = convert([w, h], box)
out_file.write(str(0) + " " + " ".join([str(round(a, 3)) for a in bb]) + '\n')
if __name__ == '__main__':
rootdir = r'D:\SoftWareInstallMenu\JetBrains\PycharmProjects\tunnel_datadeal\makesegdata\11'
out_dir = r'D:\SoftWareInstallMenu\JetBrains\PycharmProjects\tunnel_datadeal\makesegdata\11_save'
cls_list = countpixel(rootdir) # 所有类别,第零位是[0,0,0]背景不算做类别,所以cls_id要减1
convert_annotation(rootdir, out_dir, cls_list)
转多类别
import os
import cv2
# 传入图像地址,统计某张照片中的所有像素值
def countpixel(rootdir):
assert os.path.exists(rootdir), f'{rootdir} not exist.'
arr = [] # 查看了tunnelmask的像素值
for file in os.listdir(rootdir):
img = cv2.imread(os.path.join(rootdir, file))
height, width = img.shape[:-1]
for h in range(height):
for w in range(width):
pixel = img[h, w].tolist()
if pixel not in arr:
arr.append(pixel)
print(f'arr={arr}')
return arr
# 从一张图中找到[1, 1, 1]出现的x_min x_max y_min y_max
def position_lwd(imagepath, cls_list):
image = cv2.imread(imagepath)
height, width = image.shape[:-1]
val = []
x = []
y = []
for h in range(height):
for w in range(width):
pixel = image[h, w].tolist()
if pixel != [0, 0, 0]: # 如果不为背景
if pixel not in val:
val.append(pixel)
x.append([])
y.append([])
index = val.index(pixel)
x[index].append(w)
y[index].append(h)
# 找出每一个类别的x_min x_max y_min y_max
boxs = dict()
for i in range(len(val)):
cls_id = cls_list.index(val[i]) - 1 # 类别编号
boxs[cls_id] = []
boxs[cls_id].append(min(x[i]))
boxs[cls_id].append(max(x[i]))
boxs[cls_id].append(min(y[i]))
boxs[cls_id].append(max(y[i]))
return boxs
# 归一化处理
def convert(size, box):
# 该函数将xmin,ymin,xmax,ymax转为x,y,w,h中心点坐标和宽高
dw = 1. / (size[0])
dh = 1. / (size[1])
x = (box[0] + box[1]) / 2.0
y = (box[2] + box[3]) / 2.0
w = box[1] - box[0]
h = box[3] - box[2]
x = x * dw
w = w * dw
y = y * dh
h = h * dh
return (x, y, w, h)
def convert_annotation(rootdir, out_dir, cls_list):
assert os.path.exists(rootdir), f'{rootdir} not exist.'
if not os.path.exists(out_dir):
os.makedirs(out_dir)
filelist = os.listdir(rootdir)
for file in filelist:
filename = os.path.splitext(file)[0]
file_path = os.path.join(rootdir, file)
file_cv = cv2.imread(file_path)
h, w = file_cv.shape[:-1]
# 输出标签的地址
out_path = os.path.join(out_dir, filename + '.txt')
out_file = open(out_path, 'w')
# 开始转换
boxs = position_lwd(file_path, cls_list) # 多类别转换
for cls_id, box in boxs.items():
bb = convert([w, h], box)
out_file.write(str(cls_id) + " " + " ".join([str(round(a, 3)) for a in bb]) + '\n')
if __name__ == '__main__':
rootdir = r'D:\SoftWareInstallMenu\JetBrains\PycharmProjects\yolov5ds_lwd\datadeal\makesegdata\tunnel_mask'
out_dir = r'D:\SoftWareInstallMenu\JetBrains\PycharmProjects\yolov5ds_lwd\datadeal\makesegdata\tunnel_mask_out'
cls_list = countpixel(rootdir) # 所有类别,第零位是[0,0,0]背景不算做类别,所以cls_id要减1
convert_annotation(rootdir, out_dir, cls_list)