You cannot select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
71 lines
2.2 KiB
Python
71 lines
2.2 KiB
Python
1 year ago
|
from ultralytics import YOLO
|
||
|
|
from PIL import Image
|
||
|
|
import cv2
|
||
|
|
import os
|
||
|
|
|
||
|
|
# def draw_boxes(image, detections):
|
||
|
|
# for detection in detections:
|
||
|
|
# label_name, bbox_list = detection['label'], detection['bbox']
|
||
|
|
# confidence = detection['confidence']
|
||
|
|
|
||
|
|
# # 提取边界框坐标
|
||
|
|
# x1, y1, x2, y2 = map(int, bbox_list)
|
||
|
|
|
||
|
|
# # 在图像上绘制边界框
|
||
|
|
# cv2.rectangle(image, (x1, y1), (x2, y2), (0, 255, 0), 2)
|
||
|
|
|
||
|
|
# # 在图像上写入标签名和置信度
|
||
|
|
# text = f"{label_name}: {confidence}"
|
||
|
|
# cv2.putText(image, text, (x1, y1 - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)
|
||
|
|
|
||
|
|
# return image
|
||
|
|
|
||
|
|
# def analysis_yolov8(frame, model_coco, confidence_set):
|
||
|
|
# results_coco = model_coco(frame)
|
||
|
|
# re_list = []
|
||
|
|
|
||
|
|
# if results_coco:
|
||
|
|
# for r in results_coco:
|
||
|
|
# boxes = r.boxes
|
||
|
|
# for box in boxes:
|
||
|
|
# b = box.xyxy[0] # 获取边界框坐标
|
||
|
|
# c = box.cls
|
||
|
|
# labels_name = model_coco.names[int(c)]
|
||
|
|
# confidence = float(box.conf)
|
||
|
|
# confidence = round(confidence, 2)
|
||
|
|
|
||
|
|
# if confidence < confidence_set:
|
||
|
|
# continue
|
||
|
|
|
||
|
|
|
||
|
|
# re_dict = {'label': labels_name, 'bbox': b, 'confidence': confidence}
|
||
|
|
# re_list.append(re_dict)
|
||
|
|
|
||
|
|
# return re_list
|
||
|
|
|
||
|
|
# # # 创建新文件夹保存结果
|
||
|
|
# output_folder = "E:/code_files/jy/20240411_test_img_file/t_0"
|
||
|
|
# os.makedirs(output_folder, exist_ok=True)
|
||
|
|
|
||
|
|
# 加载模型和图像
|
||
|
|
model = YOLO("E:/code_files/jy/20240411_model_test/train/weights/best.onnx")
|
||
|
|
img_files = r"E:\code_files\jy\20240423"
|
||
|
|
|
||
|
|
for filename in os.listdir(img_files):
|
||
|
|
|
||
|
|
img_path = os.path.join(img_files, filename)
|
||
|
|
model.predict(img_path, save=True)
|
||
|
|
# frame = cv2.imread(img_path) # 替换为你的图像路径
|
||
|
|
|
||
|
|
# 调用检测函数获取边界框结果
|
||
|
|
# detections = analysis_yolov8(frame, model, confidence_set=0.5)
|
||
|
|
|
||
|
|
# # 绘制边界框和标签
|
||
|
|
# annotated_image = draw_boxes(frame, detections)
|
||
|
|
|
||
|
|
# # 保存结果图像到新文件夹
|
||
|
|
# output_path = os.path.join(output_folder, filename)
|
||
|
|
# cv2.imwrite(output_path, annotated_image)
|
||
|
|
|
||
|
|
# print(f"Annotated image saved at: {output_path}")
|