本文共 3964 字，大约阅读时间需要 13 分钟。

yolov5 github：https://github.com/ultralytics/yolov5

detect:

from torchvision import transformsimport torchfrom PIL import Image,ImageDrawfrom models import yolofrom utils.general import non_max_suppressionfrom models.experimental import attempt_load# model = yolo.Model(r"D:\GoogleEarthProPortable\yolov5-master\models\yolov5s.yaml")# model.load_state_dict(torch.load(r"D:\GoogleEarthProPortable\yolov5-master\weights\yolov5s.pt"))model = attempt_load("weights/yolov5s.pt")  # load FP32 modelmodel.eval()img = Image.open("inference/images/bus.jpg")tf = transforms.Compose([    transforms.Resize((512,640)),    transforms.ToTensor()])print(img.size) # w,hscale_w = img.size[0] /640scale_h = img.size[1] /512im = img.resize((640,512))img_tensor = tf(img)pred = model(img_tensor[None])[0]pred = non_max_suppression(pred,0.3,0.5)imgDraw = ImageDraw.Draw(img)for box in pred[0]:    b = box.cpu().detach().long().numpy()    print(b)    imgDraw.rectangle((b[0]*scale_w,b[1]*scale_h,b[2]*scale_w,b[3]*scale_h))    # imgDraw.rectangle((b[0],b[1],b[2],b[3]))img.show()

serving：

import ioimport jsonfrom torchvision import modelsimport torchvision.transforms as transformsfrom PIL import Image,ImageDrawfrom utils.general import non_max_suppressionfrom models.experimental import attempt_loadfrom flask import Flask, jsonify, requestapp = Flask(__name__)model = attempt_load("weights/yolov5s.pt")  # load FP32 modelmodel.eval()names= ['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 transform_image(image_bytes):    my_transforms = transforms.Compose([transforms.Resize((512,640)),                                        transforms.ToTensor(),                                        ])    image = Image.open(io.BytesIO(image_bytes))    return my_transforms(image)def get_prediction(image_bytes):    tensor = transform_image(image_bytes=image_bytes)    outputs = model(tensor[None])[0]    print(outputs)    outputs = non_max_suppression(outputs,0.3,0.5)    boxs = outputs[0]    print(boxs[0])    print(int(boxs[0][-1].item()))    class_name = names[int(boxs[0][5].item())]    print(boxs.shape)    boxes = []    for i in range(boxs.shape[0]):        boxes.append([boxs[i][0].item(),boxs[i][1].item(),boxs[i][2].item(),boxs[i][3].item(),boxs[i][4].item(),boxs[i][5].item()])    return boxes@app.route('/predict', methods=['POST'])def predict():    if request.method == 'POST':        file = request.files['file']        img_bytes = file.read()        boxes = get_prediction(image_bytes=img_bytes)        return ({'boxes': boxes})if __name__ == '__main__':    app.run()

client：

import requestsimport osfor i in os.listdir("inference/images"):    image = open("inference/images/"+i,'rb')    payload = {'file':image}    r = requests.post(" http://localhost:5000/predict", files=payload).json()    print(r)

git bash控制台：