好的,你想用训练好的YOLO模型处理图片,并且只保留带框的部分。这个需求可以分为以下几个关键步骤:
1. **加载训练好的YOLO模型**:确保模型路径正确。
2. **用模型进行预测**:检测目标并返回包含的边界框(bounding box)。
3. **从输入图片中裁剪出包含目标的区域**:只保留带框的部分。
我们来一步步实现这个过程。
### 1. 加载训练好的YOLO模型
首先,确保路径正确,再进行加载。
import torch
from PIL import Image
import cv2
# 确保你有一个有效的模型路径
model_path = 'C:\\Users\\shiro\\Desktop\\yolov5-mask-42-master\\runs\\train\\exp19\\last.pt'
model = torch.load(model_path, map_location=torch.device('cpu')) # 如果没有GPU,使用CPU
### 2. 用模型进行预测
接下来,用YOLO模型对图片进行预测。
# 加载图片
image_path = 'path_to_your_image.jpg'
img = Image.open(image_path)
# 进行预测
results = model(img)
# 获取预测结果
predictions = results.xyxy[0] # 选择第一个图像的预测结果
### 3. 从输入图片中裁剪出包含目标的区域
遍历预测结果中的每个边界框, 并将其裁剪出来。
# 转换图像为OpenCV格式 (YOLOv5默认使用PIL图像)
img_cv = cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR)
cropped_images = []
for pred in predictions:
x1, y1, x2, y2, confidence, class_id = pred
x1, y1, x2, y2 = int(x1), int(y1), int(x2), int(y2)
# 裁剪带框部分
cropped_image = img_cv[y1:y2, x1:x2]
cropped_images.append(cropped_image)
# 保存裁剪后的图片 (可选)
cropped_image_path = f'cropped_{x1}_{y1}_{x2}_{y2}.jpg'
cv2.imwrite(cropped_image_path, cropped_image)
注意,这段代码假设你的YOLO模型的输出格式是 `[x1, y1, x2, y2, confidence, class_id]`,你需要根据实际情况进行调整。
### 4. 合并所有步骤
最后,把所有步骤合并到一起。
import torch
from PIL import Image
import cv2
import numpy as np
def load_model(model_path):
return torch.load(model_path, map_location=torch.device('cpu'))
def predict_and_crop(model, image_path):
img = Image.open(image_path)
results = model(img)
predictions = results.xyxy[0]
img_cv = cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR)
cropped_images = []
for pred in predictions:
x1, y1, x2, y2, confidence, class_id = pred
x1, y1, x2, y2 = int(x1), int(y1), int(x2), int(y2)
cropped_image = img_cv[y1:y2, x1:x2]
cropped_images.append(cropped_image)
cropped_image_path = f'cropped_{x1}_{y1}_{x2}_{y2}.jpg'
cv2.imwrite(cropped_image_path, cropped_image)
return cropped_images
model_path = 'C:\\Users\\shiro\\Desktop\\yolov5-mask-42-master\\runs\\train\\exp19\\last.pt'
image_path = 'path_to_your_image.jpg'
model = load_model(model_path)
cropped_images = predict_and_crop(model, image_path)
### 总结
以上代码实现了从加载训练好的YOLO模型,对图片进行目标检测,并且只保留带框部分的功能。你只需更改 `model_path` 和 `image_path` 为你实际的路径。希望对你有所帮助,祝你愉快使用YOLO模型!如果遇到任何问题,请随时问我。
以上回复来自 FishC AI 小助理,如未能正确解答您的问题,请继续追问。 |