290 lines
8.0 KiB
Python
290 lines
8.0 KiB
Python
# 该脚本用来加载模型并执行一张图像的前向推理
|
|
import sys
|
|
sys.path.append(R"../0_yolov5_7.0")
|
|
import argparse
|
|
import os
|
|
import cv2
|
|
import torch
|
|
import numpy as np
|
|
from utils.augmentations import letterbox
|
|
from utils.general import check_version, non_max_suppression
|
|
from models.experimental import attempt_load
|
|
|
|
# 画图所需数值
|
|
_COLORS = np.array(
|
|
[
|
|
0.000, 0.447, 0.741,
|
|
0.850, 0.325, 0.098,
|
|
0.929, 0.694, 0.125,
|
|
0.494, 0.184, 0.556,
|
|
0.466, 0.674, 0.188,
|
|
0.301, 0.745, 0.933,
|
|
0.635, 0.078, 0.184,
|
|
0.300, 0.300, 0.300,
|
|
0.600, 0.600, 0.600,
|
|
1.000, 0.000, 0.000,
|
|
1.000, 0.500, 0.000,
|
|
0.749, 0.749, 0.000,
|
|
0.000, 1.000, 0.000,
|
|
0.000, 0.000, 1.000,
|
|
0.667, 0.000, 1.000,
|
|
0.333, 0.333, 0.000,
|
|
0.333, 0.667, 0.000,
|
|
0.333, 1.000, 0.000,
|
|
0.667, 0.333, 0.000,
|
|
0.667, 0.667, 0.000,
|
|
0.667, 1.000, 0.000,
|
|
1.000, 0.333, 0.000,
|
|
1.000, 0.667, 0.000,
|
|
1.000, 1.000, 0.000,
|
|
0.000, 0.333, 0.500,
|
|
0.000, 0.667, 0.500,
|
|
0.000, 1.000, 0.500,
|
|
0.333, 0.000, 0.500,
|
|
0.333, 0.333, 0.500,
|
|
0.333, 0.667, 0.500,
|
|
0.333, 1.000, 0.500,
|
|
0.667, 0.000, 0.500,
|
|
0.667, 0.333, 0.500,
|
|
0.667, 0.667, 0.500,
|
|
0.667, 1.000, 0.500,
|
|
1.000, 0.000, 0.500,
|
|
1.000, 0.333, 0.500,
|
|
1.000, 0.667, 0.500,
|
|
1.000, 1.000, 0.500,
|
|
0.000, 0.333, 1.000,
|
|
0.000, 0.667, 1.000,
|
|
0.000, 1.000, 1.000,
|
|
0.333, 0.000, 1.000,
|
|
0.333, 0.333, 1.000,
|
|
0.333, 0.667, 1.000,
|
|
0.333, 1.000, 1.000,
|
|
0.667, 0.000, 1.000,
|
|
0.667, 0.333, 1.000,
|
|
0.667, 0.667, 1.000,
|
|
0.667, 1.000, 1.000,
|
|
1.000, 0.000, 1.000,
|
|
1.000, 0.333, 1.000,
|
|
1.000, 0.667, 1.000,
|
|
0.333, 0.000, 0.000,
|
|
0.500, 0.000, 0.000,
|
|
0.667, 0.000, 0.000,
|
|
0.833, 0.000, 0.000,
|
|
1.000, 0.000, 0.000,
|
|
0.000, 0.167, 0.000,
|
|
0.000, 0.333, 0.000,
|
|
0.000, 0.500, 0.000,
|
|
0.000, 0.667, 0.000,
|
|
0.000, 0.833, 0.000,
|
|
0.000, 1.000, 0.000,
|
|
0.000, 0.000, 0.167,
|
|
0.000, 0.000, 0.333,
|
|
0.000, 0.000, 0.500,
|
|
0.000, 0.000, 0.667,
|
|
0.000, 0.000, 0.833,
|
|
0.000, 0.000, 1.000,
|
|
0.000, 0.000, 0.000,
|
|
0.143, 0.143, 0.143,
|
|
0.286, 0.286, 0.286,
|
|
0.429, 0.429, 0.429,
|
|
0.571, 0.571, 0.571,
|
|
0.714, 0.714, 0.714,
|
|
0.857, 0.857, 0.857,
|
|
0.000, 0.447, 0.741,
|
|
0.314, 0.717, 0.741,
|
|
0.50, 0.5, 0
|
|
]
|
|
).astype(np.float32).reshape(-1, 3)
|
|
|
|
COCO_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",
|
|
)
|
|
|
|
|
|
def vis(img, boxes, scores, cls_ids, conf=0.5, class_names=None):
|
|
|
|
for i in range(len(boxes)):
|
|
box = boxes[i]
|
|
cls_id = int(cls_ids[i])
|
|
score = scores[i]
|
|
if score < conf:
|
|
continue
|
|
x0 = int(box[0])
|
|
y0 = int(box[1])
|
|
x1 = int(box[2])
|
|
y1 = int(box[3])
|
|
|
|
color = (_COLORS[cls_id] * 255).astype(np.uint8).tolist()
|
|
text = '{}:{:.1f}%'.format(class_names[cls_id], score * 100)
|
|
txt_color = (0, 0, 0) if np.mean(_COLORS[cls_id]) > 0.5 else (255, 255, 255)
|
|
font = cv2.FONT_HERSHEY_SIMPLEX
|
|
|
|
txt_size = cv2.getTextSize(text, font, 0.4, 1)[0]
|
|
cv2.rectangle(img, (x0, y0), (x1, y1), color, 2)
|
|
|
|
txt_bk_color = (_COLORS[cls_id] * 255 * 0.7).astype(np.uint8).tolist()
|
|
cv2.rectangle(
|
|
img,
|
|
(x0, y0 + 1),
|
|
(x0 + txt_size[0] + 1, y0 + int(1.5*txt_size[1])),
|
|
txt_bk_color,
|
|
-1
|
|
)
|
|
cv2.putText(img, text, (x0, y0 + txt_size[1]), font, 0.4, txt_color, thickness=1)
|
|
|
|
return img
|
|
|
|
|
|
def scale_coords(img1_shape, coords, img0_shape, ratio_pad=None):
|
|
# Rescale coords (xyxy) from img1_shape to img0_shape
|
|
if ratio_pad is None: # calculate from img0_shape
|
|
gain = min(img1_shape[0] / img0_shape[0], img1_shape[1] / img0_shape[1]) # gain = old / new
|
|
pad = (img1_shape[1] - img0_shape[1] * gain) / 2, (img1_shape[0] - img0_shape[0] * gain) / 2 # wh padding
|
|
else:
|
|
gain = ratio_pad[0][0]
|
|
pad = ratio_pad[1]
|
|
|
|
coords[:, [0, 2]] -= pad[0] # x padding
|
|
coords[:, [1, 3]] -= pad[1] # y padding
|
|
coords[:, :4] /= gain
|
|
clip_coords(coords, img0_shape)
|
|
return coords
|
|
|
|
|
|
def clip_coords(boxes, shape):
|
|
# Clip bounding xyxy bounding boxes to image shape (height, width)
|
|
if isinstance(boxes, torch.Tensor): # faster individually
|
|
boxes[:, 0].clamp_(0, shape[1]) # x1
|
|
boxes[:, 1].clamp_(0, shape[0]) # y1
|
|
boxes[:, 2].clamp_(0, shape[1]) # x2
|
|
boxes[:, 3].clamp_(0, shape[0]) # y2
|
|
else: # np.array (faster grouped)
|
|
boxes[:, [0, 2]] = boxes[:, [0, 2]].clip(0, shape[1]) # x1, x2
|
|
boxes[:, [1, 3]] = boxes[:, [1, 3]].clip(0, shape[0]) # y1, y2
|
|
|
|
|
|
def pred_one_image(img_path, model_path, test_size):
|
|
|
|
img_raw = cv2.imread(img_path)
|
|
|
|
# 前处理
|
|
img = letterbox(img_raw, new_shape=test_size, stride=32, auto=True)[0] # auto参数设为True表示带pad的resize
|
|
img = img.transpose((2, 0, 1))[::-1] # HWC to CHW, BGR to RGB
|
|
img = np.ascontiguousarray(img)
|
|
im = torch.from_numpy(img).float().unsqueeze(0)
|
|
im /= 255
|
|
|
|
# 加载模型
|
|
model = attempt_load(model_path)
|
|
pred = model(im)
|
|
|
|
# NMS
|
|
conf_thres = 0.25
|
|
iou_thres = 0.45
|
|
pred = non_max_suppression(pred, conf_thres, iou_thres, classes=None, agnostic=False, max_det=1000)
|
|
|
|
# 结果显示
|
|
det = pred[0]
|
|
det[:, :4] = scale_coords(im.shape[2:], det[:, :4], img_raw.shape).round()
|
|
result_image = vis(img_raw, boxes=det[:,:4], scores=det[:,4], cls_ids=det[:,5], conf=conf_thres, class_names=COCO_CLASSES)
|
|
cv2.imshow(" ", result_image)
|
|
cv2.waitKey(0)
|
|
|
|
if __name__ == "__main__":
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument('--model_pt', type=str, default=R'../weights/yolov5s.pt', help='model path')
|
|
parser.add_argument('--source', type=str, default=R'../0_yolov5_7.0/data/images/bus.jpg', help='image path')
|
|
parser.add_argument('--imgsz', nargs='+', type=int, default=[640], help='image size')
|
|
opt = parser.parse_args()
|
|
|
|
opt.imgsz *= 2 if len(opt.imgsz) == 1 else 1
|
|
test_size = tuple(opt.imgsz)
|
|
|
|
if os.path.isfile(opt.source):
|
|
pred_one_image(opt.source, opt.model_pt, test_size)
|
|
elif os.path.isdir(opt.source):
|
|
image_list = os.listdir(opt.source)
|
|
for image_file in image_list:
|
|
image_path = opt.source + "//" + image_file
|
|
pred_one_image(image_path, opt.model_pt, test_size)
|
|
|
|
|