第一次提交Yolo项目
This commit is contained in:
lhr
189
inference.py
Normal file
189
inference.py
Normal file
@@ -0,0 +1,189 @@
|
||||
import torch
|
||||
import cv2
|
||||
import numpy as np
|
||||
import torchvision
|
||||
from yolo11_standalone import YOLO11
|
||||
|
||||
# COCO 80类 类别名称
|
||||
CLASSES = [
|
||||
"person", "bicycle", "car", "motorcycle", "airplane", "bus", "train", "truck", "boat", "traffic light",
|
||||
"fire hydrant", "stop sign", "parking meter", "bench", "bird", "cat", "dog", "horse", "sheep", "cow",
|
||||
"elephant", "bear", "zebra", "giraffe", "backpack", "umbrella", "handbag", "tie", "suitcase", "frisbee",
|
||||
"skis", "snowboard", "sports ball", "kite", "baseball bat", "baseball glove", "skateboard", "surfboard",
|
||||
"tennis racket", "bottle", "wine glass", "cup", "fork", "knife", "spoon", "bowl", "banana", "apple",
|
||||
"sandwich", "orange", "broccoli", "carrot", "hot dog", "pizza", "donut", "cake", "chair", "couch",
|
||||
"potted plant", "bed", "dining table", "toilet", "tv", "laptop", "mouse", "remote", "keyboard", "cell phone",
|
||||
"microwave", "oven", "toaster", "sink", "refrigerator", "book", "clock", "vase", "scissors", "teddy bear",
|
||||
"hair drier", "toothbrush"
|
||||
]
|
||||
|
||||
# 生成随机颜色用于绘图
|
||||
COLORS = np.random.uniform(0, 255, size=(len(CLASSES), 3))
|
||||
|
||||
def letterbox(im, new_shape=(640, 640), color=(114, 114, 114)):
|
||||
"""
|
||||
将图像缩放并填充到指定大小 (保持纵横比)
|
||||
"""
|
||||
shape = im.shape[:2] # current shape [height, width]
|
||||
if isinstance(new_shape, int):
|
||||
new_shape = (new_shape, new_shape)
|
||||
|
||||
# 计算缩放比例
|
||||
r = min(new_shape[0] / shape[0], new_shape[1] / shape[1])
|
||||
|
||||
# 计算padding
|
||||
new_unpad = int(round(shape[1] * r)), int(round(shape[0] * r))
|
||||
dw, dh = new_shape[1] - new_unpad[0], new_shape[0] - new_unpad[1] # wh padding
|
||||
dw, dh = dw / 2, dh / 2 # divide padding into 2 sides
|
||||
|
||||
if shape[::-1] != new_unpad: # resize
|
||||
im = cv2.resize(im, new_unpad, interpolation=cv2.INTER_LINEAR)
|
||||
|
||||
top, bottom = int(round(dh - 0.1)), int(round(dh + 0.1))
|
||||
left, right = int(round(dw - 0.1)), int(round(dw + 0.1))
|
||||
|
||||
# 添加边框
|
||||
im = cv2.copyMakeBorder(im, top, bottom, left, right, cv2.BORDER_CONSTANT, value=color)
|
||||
return im, r, (dw, dh)
|
||||
|
||||
def xywh2xyxy(x):
|
||||
"""Convert nx4 boxes from [x, y, w, h] to [x1, y1, x2, y2]"""
|
||||
y = x.clone() if isinstance(x, torch.Tensor) else np.copy(x)
|
||||
y[..., 0] = x[..., 0] - x[..., 2] / 2 # top left x
|
||||
y[..., 1] = x[..., 1] - x[..., 3] / 2 # top left y
|
||||
y[..., 2] = x[..., 0] + x[..., 2] / 2 # bottom right x
|
||||
y[..., 3] = x[..., 1] + x[..., 3] / 2 # bottom right y
|
||||
return y
|
||||
|
||||
def non_max_suppression(prediction, conf_thres=0.25, iou_thres=0.45, max_det=300):
|
||||
"""
|
||||
非极大值抑制 (NMS)
|
||||
prediction: [Batch, 84, Anchors]
|
||||
"""
|
||||
# 1. 转置: [Batch, 84, Anchors] -> [Batch, Anchors, 84]
|
||||
prediction = prediction.transpose(1, 2)
|
||||
|
||||
bs = prediction.shape[0] # batch size
|
||||
nc = prediction.shape[2] - 4 # number of classes
|
||||
|
||||
# 修复: 使用 max(-1) 在最后一个维度(类别)上寻找最大置信度
|
||||
# 之前的 max(1) 错误地在 Anchors 维度上操作了
|
||||
xc = prediction[..., 4:].max(-1)[0] > conf_thres # candidates
|
||||
|
||||
output = [torch.zeros((0, 6), device=prediction.device)] * bs
|
||||
|
||||
for xi, x in enumerate(prediction): # image index, image inference
|
||||
x = x[xc[xi]] # confidence filtering
|
||||
|
||||
if not x.shape[0]:
|
||||
continue
|
||||
|
||||
# Box (center x, center y, width, height) to (x1, y1, x2, y2)
|
||||
box = xywh2xyxy(x[:, :4])
|
||||
|
||||
# Confidence and Class
|
||||
conf, j = x[:, 4:].max(1, keepdim=True)
|
||||
x = torch.cat((box, conf, j.float()), 1)[conf.view(-1) > conf_thres]
|
||||
|
||||
# Check shape
|
||||
n = x.shape[0]
|
||||
if not n:
|
||||
continue
|
||||
elif n > max_det:
|
||||
x = x[x[:, 4].argsort(descending=True)[:max_det]]
|
||||
|
||||
# Batched NMS
|
||||
c = x[:, 5:6] * 7680 # classes
|
||||
boxes, scores = x[:, :4] + c, x[:, 4]
|
||||
i = torchvision.ops.nms(boxes, scores, iou_thres)
|
||||
output[xi] = x[i]
|
||||
|
||||
return output
|
||||
|
||||
def main():
|
||||
# 1. 初始化模型
|
||||
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
||||
print(f"Using device: {device}")
|
||||
|
||||
model = YOLO11(nc=80, scale='s')
|
||||
# 加载你之前转换好的纯净权重
|
||||
model.load_weights("yolo11s.pth")
|
||||
model.to(device)
|
||||
model.eval()
|
||||
# model.train()
|
||||
|
||||
# 2. 读取图片
|
||||
img_path = "1.jpg" # 请替换为你本地的图片路径
|
||||
|
||||
img0 = cv2.imread(img_path)
|
||||
assert img0 is not None, f"Image Not Found {img_path}"
|
||||
|
||||
# 3. 预处理
|
||||
# Letterbox resize
|
||||
img, ratio, (dw, dh) = letterbox(img0, new_shape=(640, 640))
|
||||
|
||||
# BGR to RGB, HWC to CHW
|
||||
img = img[:, :, ::-1].transpose(2, 0, 1)
|
||||
img = np.ascontiguousarray(img)
|
||||
|
||||
img_tensor = torch.from_numpy(img).to(device)
|
||||
img_tensor = img_tensor.float()
|
||||
img_tensor /= 255.0 # 0 - 255 to 0.0 - 1.0
|
||||
if img_tensor.ndim == 3:
|
||||
img_tensor = img_tensor.unsqueeze(0)
|
||||
|
||||
# 4. 推理
|
||||
print("开始推理...")
|
||||
with torch.no_grad():
|
||||
pred = model(img_tensor)
|
||||
|
||||
# 5. 后处理 (NMS)
|
||||
pred = non_max_suppression(pred, conf_thres=0.25, iou_thres=0.45)
|
||||
|
||||
# 6. 绘制结果
|
||||
det = pred[0] # 仅处理第一张图片
|
||||
|
||||
if len(det):
|
||||
# 将坐标映射回原图尺寸
|
||||
# det[:, :4] 是 x1, y1, x2, y2
|
||||
det[:, [0, 2]] -= dw # x padding
|
||||
det[:, [1, 3]] -= dh # y padding
|
||||
det[:, :4] /= ratio
|
||||
|
||||
# 裁剪坐标防止越界
|
||||
det[:, 0].clamp_(0, img0.shape[1])
|
||||
det[:, 1].clamp_(0, img0.shape[0])
|
||||
det[:, 2].clamp_(0, img0.shape[1])
|
||||
det[:, 3].clamp_(0, img0.shape[0])
|
||||
|
||||
print(f"检测到 {len(det)} 个目标")
|
||||
|
||||
for *xyxy, conf, cls in det:
|
||||
c = int(cls)
|
||||
label = f'{CLASSES[c]} {conf:.2f}'
|
||||
p1, p2 = (int(xyxy[0]), int(xyxy[1])), (int(xyxy[2]), int(xyxy[3]))
|
||||
|
||||
# 画框
|
||||
color = COLORS[c]
|
||||
cv2.rectangle(img0, p1, p2, color, 2, lineType=cv2.LINE_AA)
|
||||
|
||||
# 画标签背景
|
||||
t_size = cv2.getTextSize(label, 0, fontScale=0.5, thickness=1)[0]
|
||||
p2_label = p1[0] + t_size[0], p1[1] - t_size[1] - 3
|
||||
cv2.rectangle(img0, p1, p2_label, color, -1, cv2.LINE_AA)
|
||||
|
||||
# 画文字
|
||||
cv2.putText(img0, label, (p1[0], p1[1] - 2), 0, 0.5, [255, 255, 255], thickness=1, lineType=cv2.LINE_AA)
|
||||
|
||||
print(f" - {label} at {p1}-{p2}")
|
||||
|
||||
# 7. 显示/保存结果
|
||||
cv2.imwrite("result.jpg", img0)
|
||||
print("结果已保存至 result.jpg")
|
||||
|
||||
def import_os_exists(path):
|
||||
import os
|
||||
return os.path.exists(path)
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user