YuXin_Liu/丝束检测/drawGroundTruth.py

68 lines
2.1 KiB
Python

import cv2
import os
def draw_boxes(image, labels, w, h):
for line in labels:
parts = line.strip().split()
class_id = int(parts[0])
x_center = float(parts[1])
y_center = float(parts[2])
width = float(parts[3])
height = float(parts[4])
# 计算边框的左上角和右下角坐标
x_center *= w
y_center *= h
width *= w
height *= h
x_min = int(x_center - width / 2)
y_min = int(y_center - height / 2)
x_max = int(x_center + width / 2)
y_max = int(y_center + height / 2)
# 画出边框
color = (0, 255, 0) # 绿色边框
thickness = 10
image = cv2.rectangle(image, (x_min, y_min), (x_max, y_max), color, thickness)
return image
def process_folder(image_folder, label_folder, output_folder):
if not os.path.exists(output_folder):
os.makedirs(output_folder)
image_files = [f for f in os.listdir(image_folder) if f.endswith(('.jpg', '.jpeg', '.png'))]
for image_file in image_files:
image_path = os.path.join(image_folder, image_file)
label_file = os.path.splitext(image_file)[0] + '.txt'
label_path = os.path.join(label_folder, label_file)
if os.path.exists(label_path):
# 读取图片
image = cv2.imread(image_path)
h, w, _ = image.shape
# 读取标签文件
with open(label_path, 'r') as file:
labels = file.readlines()
# 绘制边框
image_with_boxes = draw_boxes(image, labels, w, h)
# 保存结果图片
output_path = os.path.join(output_folder, image_file)
cv2.imwrite(output_path, image_with_boxes)
print(f'Processed {image_file}')
else:
print(f'No label file found for {image_file}')
# 示例用法
image_folder = 'testTiaoJuan/cutJPGOnlyTiaoJuan'
label_folder = 'testTiaoJuan/CutedLabels_tiaojuan'
output_folder = 'runs/test/cuted_conf0.035'
process_folder(image_folder, label_folder, output_folder)