39 lines
1.1 KiB
Python
39 lines
1.1 KiB
Python
import os
|
||
import cv2
|
||
from ultralytics import YOLO
|
||
|
||
# 加载模型
|
||
model = YOLO("best.pt")
|
||
|
||
# 输入文件夹路径
|
||
input_dir = "test_images/"
|
||
temp_dir = "temp_rgb_images"
|
||
|
||
# 创建临时文件夹用于保存转换后的RGB图片
|
||
os.makedirs(temp_dir, exist_ok=True)
|
||
|
||
# 扫描文件夹，将灰度图转换为RGB
|
||
for file in os.listdir(input_dir):
|
||
if file.lower().endswith(('.jpg', '.jpeg', '.png', '.bmp', '.tiff')):
|
||
path = os.path.join(input_dir, file)
|
||
img = cv2.imread(path, cv2.IMREAD_UNCHANGED)
|
||
|
||
# 如果是灰度图，转换为3通道
|
||
if len(img.shape) == 2 or img.shape[2] == 1:
|
||
img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)
|
||
|
||
# 保存到临时目录
|
||
cv2.imwrite(os.path.join(temp_dir, file), img)
|
||
|
||
# 使用转换后的图像文件夹进行检测
|
||
results = model.predict(source=temp_dir, conf=0.2, save=True)
|
||
|
||
# 输出检测信息
|
||
for result in results:
|
||
boxes = result.boxes # 检测框
|
||
for box in boxes:
|
||
cls = int(box.cls[0])
|
||
conf = float(box.conf[0])
|
||
xyxy = box.xyxy[0].tolist()
|
||
print(f"类别: {model.names[cls]}, 置信度: {conf:.2f}, 坐标: {xyxy}")
|