使用Yolov5进行目标检测，比对txt标签，标记不合格产品。
以上回复有部分引自免费微信小程序：皆我百晓生。 如问题尚未解决，还请稍待，如有疑问请回复进一步沟通。
为了根据TXT标签文件中的信息来判断产品是否合格，并在相应的图片上标注产品和不合格结果，你需要进行以下几个步骤：
- 安装依赖库： 确保已经安装了PyTorch、YOLOv5和其他相关库。如果没有安装，可以使用以下命令安装：
pip install torch torchvision
git clone https://github.com/ultralytics/yolov5.git
cd yolov5
pip install -e .
-
准备数据： 确保你有一个包含图片和对应TXT标签文件的数据集。TXT文件应该有两列，第一列是图片名称（不带扩展名），第二列是标签。
-
编写Python脚本： 创建一个名为
label_images.py的文件，然后添加以下代码：
import os
import cv2
from PIL import Image, ImageDraw, ImageFont
from typing import List
from yolov5.utils.datasets import LoadImages, LoadStreams
from yolov5.models.common import DetectMultiBackend
from yolov5.utils.general import check_img_size, non_max_suppression, scale_coords, xyxy2xywh, plot_one_box
def draw_boxes(image_path: str, labels: List[str], scores: List[float]):
img = cv2.imread(image_path)
height, width = img.shape[:2]
# Create a blank image for drawing the bounding boxes
draw_image = Image.fromarray(img)
draw = ImageDraw.Draw(draw_image)
font = ImageFont.truetype("arial.ttf", 14)
for label, score in zip(labels, scores):
if score < 0.5: # 可以根据需要调整阈值
text = f"{label} (不合格)"
else:
text = f"{label} (合格)"
# Convert YOLO's xyxy format to OpenCV's xywh format
box = [int(coord * width) for coord in scale_coords((width, height), (0, 0, 1, 1), (labels[0].shape[1], labels[0].shape[2]))]
x, y, w, h = box
# Draw the rectangle and text
draw.rectangle(box, outline=(0, 255, 0) if score >= 0.5 else (0, 0, 255))
draw.text((x, y - 5), text, fill=(0, 255, 0) if score >= 0.5 else (0, 0, 255), font=font)
result = np.array(draw_image)
cv2.imwrite(os.path.splitext(image_path)[0] + "_labeled.jpg", result)
if __name__ == "__main__":
data_dir = "path/to/your/data/directory" # 替换为你的数据集路径
txt_files = [os.path.join(data_dir, f) for f in os.listdir(data_dir) if f.endswith(".txt")]
for txt_file in txt_files:
with open(txt_file, "r") as file:
lines = file.readlines()
image_names = []
labels = []
for line in lines:
parts = line.strip().split()
image_names.append(parts[0])
labels.append(parts[1])
model = DetectMultiBackend(weights="yolov5s.pt", device="cuda")
imgsz = check_img_size(None, s=model.stride.max()) # Use default size
for i, (path, img) in enumerate(LoadImages(txt_files=image_names, img_size=imgsz)):
outputs = model(img, augment=False, visualize=True)[0]
boxes, scores, classes, nums = non_max_suppression(outputs, conf_thres=0.5, iou_thres=0.5)
draw_boxes(path, classes, scores)
- 运行脚本： 在终端中运行以下命令：
python label_images.py
这将遍历数据集中所有TXT文件，根据标签判断产品是否合格，并在原图上标注产品及其是否合格。标注后的图片将以原始图片名称加上"_labeled.jpg"的形式保存到同一目录下。
注意：这个示例脚本假设TXT文件中每个标签都是有效的，并且图像和TXT文件之间是一对一的关系。如果你的数据结构不同，可能需要修改脚本来适应你的需求。同时，YOLOv5模型默认使用的是s.pt权重文件，你可以根据需要更改模型权重。
赞
回复
