yolov5 github:https://github.com/ultralytics/yolov5
跟踪:https://github.com/mikel-brostrom/Yolov5_DeepSort_Pytorch
TensorRT:https://github.com/TrojanXu/yolov5-tensorrt
NCNN:https://github.com/WZTENG/YOLOv5_NCNN
detect:
from torchvision import transforms
import torch
from PIL import Image,ImageDraw
from models import yolo
from utils.general import non_max_suppression
from 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 model
model.eval()img = Image.open("inference/images/bus.jpg")tf = transforms.Compose([transforms.Resize((512,640)),transforms.ToTensor()
])print(img.size) # w,h
scale_w = img.size[0] /640
scale_h = img.size[1] /512
im = 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 io
import jsonfrom torchvision import models
import torchvision.transforms as transforms
from PIL import Image,ImageDrawfrom utils.general import non_max_suppression
from models.experimental import attempt_loadfrom flask import Flask, jsonify, request
app = Flask(__name__)model = attempt_load("weights/yolov5s.pt") # load FP32 model
model.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 requests
import 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控制台:
启动flask服务器:FLASK_ENV=development FLASK_APP=app.py flask run
测试命令:curl -X POST -F file=@test_img/dog.jpg http://localhost:5000/predict