YOLO is the abbreviation of You Only Look Once. It has the extraordinary ability to identify objects in images and is often used in daily applications. So in this article, we will introduce how to use FastAPI to integrate YOLOv5, so that we can use YOLOv5 as an API to provide services to the outside world.
There are several web frameworks for Python, the two most prominent of which are Flask and FastAPI. Flask is a lightweight and flexible framework ideal for small to medium-sized web applications, while FastAPI is a relatively new framework designed to build APIs quickly and easily.
FastAPI
FastAPI is a modern web framework that can be used to build APIs using Python 3.6+ based on the OpenAPI standard. The code provided by FastAPI can be used in production environments with minimal adjustments. FastAPI’s integration with ASGI (Asynchronous Server Gateway Interface) makes it one of the fastest Python frameworks.
Set up the environment
Setting up a virtual environment is a crucial first step in software development and data science. While not always mandatory, it is highly recommended. A virtual environment isolates a project’s dependencies, ensuring stability and compatibility. This practice simplifies version management and makes projects more organized and maintainable. In short, it is a best practice for improving project quality.
python3 -m venv my_project
Then use the command activate to start the virtual environment
my_project/bin/activate
Train custom YOLOv5 model
We will not introduce the detailed training process here. First, because the official code can be used directly, and second, because we mainly introduce the steps of deployment.
Here we use 6 food datasets for training, which are pre-annotated on Roboflow to classify chicken wings, French fries, dumplings, hot dogs, donuts, waffles images. If you have your own images and want to create a dataset, you can also annotate them manually on Roboflow.
python train.py - img 416 - batch 16 - epochs 100 - data {dataset.location}/data.yaml - cfg ./models/custom_yolov5s.yaml - weights '' - name yolov5s_results - cache
You only need the above sentence to train, and the results are as follows
Accuracy§: 0.763 (average of all categories)
Recall?: 0.662 (average of all categories)
mAP50 (average accuracy at IoU 0.5): 0.712 (average across all categories)
mAP50-95: 0.4 (average of all categories)
Use FastAPI for model inference
The following is the simplest project directory structure template:
├── README.md ├── main.py # The main program ├── model │ ├── __init__.py │ └── best.pt # Custom training model ├── requirements.in # Used to generate requirements.txt ├── requirements.txt # Install dependencies: pip install -r requirements.txt ├── segmentation.py # Importing the local YOLOv5 and resize images └── yolov5 # You can obtain this from https://github.com/ultralytics/yolov5
main.py
This is a function with 3 functions:
@app.get(‘ /notify/v1/health ‘), this is the endpoint to check the health of the application. It returns a JSON response with the message “OK”. This endpoint is used to check Kubernetes’ readinessProbe and liveesprobe.
@app.post(” /object-to-json “), this endpoint handles the detection of objects in the image and returns the results in JSON format. It requires an image file as input.
@app.post(” /object-to-img “), this endpoint performs object detection and returns an image with the annotated object. It also takes an image file as input and sends the annotated image as response.
#import statement
fromfastapiimportFastAPI, File#building the API
fromsegmentationimportget_yolov5, get_image_from_bytes#segmentation functions
fromstarlette.responsesimportResponse#handling API responses
importio
fromPILimportImage
importjson
fromfastapi.middleware.corsimportCORSMiddleware
#initialize and obtain the model
model=get_yolov5()
#FastAPI application setup
app=FastAPI(
title="Custom YOLOV5 Machine Learning API",
description="""Obtain object value out of image
and return image and json result""",
version="0.0.1",
)
#CORS (Cross-Origin Resource Sharing) middleware, allows the API to be accessed from different domains or origins.
origins= [
"http://localhost",
"http://localhost:8000",
"*"
]
app.add_middleware(
CORSMiddleware,
allow_origins=origins,
allow_credentials=True,
allow_methods=["*"],
allow_headers=["*"],
)
#Checking health of application, returns “OK” JSON response
@app.get('/notify/v1/health')
defget_health():
"""
Usage on K8S
readinessProbe:
httpGet: path: /notify/v1/health
port: 80
livenessProbe:
httpGet:
path: /notify/v1/health
port: 80
:return:
dict(msg='OK')
"""
returndict(msg='OK')
@app.post("/object-to-json")
asyncdefdetect_food_return_json_result(file: bytes=File(...)):
input_image=get_image_from_bytes(file)
results=model(input_image)
detect_res=results.pandas().xyxy[0].to_json(orient="records") # JSON img1 predictions
detect_res=json.loads(detect_res)
return {"result": detect_res}
@app.post("/object-to-img")
asyncdefdetect_food_return_base64_img(file: bytes=File(...)):
input_image=get_image_from_bytes(file)
results=model(input_image)
results.render() # updates results.imgs with boxes and labels
forimginresults.imgs:
bytes_io=io.BytesIO()
img_base64=Image.fromarray(img)
img_base64.save(bytes_io, format="jpeg")
returnResponse(content=bytes_io.getvalue(), media_type="image/jpeg")
segmentation.py defines 2 methods:
1. Get_yolov5(): This is where yolov5 can use customized models.
Model = torch.hub.load(‘./yolov5’, ‘custom’, path=’./model/best.pt’, source=’local’): It loads a custom yolov5 model from the local directory. The ‘custom’ parameter specifies the model architecture, ‘./model/best.pt’ is the path to the customized training model file, and ‘source’ indicates that the model is located locally. conf = 0.5: Set the confidence threshold for object detection. Only objects with a confidence level greater than 0.5 will be included in the detection results.
2. Get_image_from_bytes(): Adjust the image size to ensure compatibility with model expectations.
input_image = image .open(io.BytesIO(binary_image)).convert(” RGB “): It reads binary image data and converts it into a PIL (Python Imaging Library) image in RGB format. It calculates the width and height of the image and determines the resizing factor to fit the specified max_size. Resized_image is created by resizing the input image according to the calculated dimensions, ensuring that it does not exceed the maximum size.
Start project
Enter the following command on the terminal
uvicorn main:app - reload - host 0.0.0.0 - port 8000
When you run the program successfully it will look like this
Test results
Open http://127.0.0.1:8000/docs#/ in your browser. You should see the Open API documentation with 3 endpoints:
/notify/v1/health returns a JSON response with the message ‘OK’. You can check this by clicking try it out and executing, the output is as follows:
/object-to-json handles the detection of objects in images and returns the results in JSON format. It requires an image file as input.
/object-to-img Performs object detection and returns an image with annotated objects. It also takes an image file as input and sends the annotated image as response.
This completes our program
Summary
Here we see the complete process of deploying the yolov5 model, and this process can be applied to the production environment. But here we found a new command uvicorn’.
Uvicorn is a very fast ASGI server built on uvloop and httptools. Because Uvicorn is a very fast, reliable and easy-to-use ASGI server that remains stable and efficient when handling large numbers of concurrent connections. Uvicorn also supports new features such as WebSocket and HTTP/2, which is in line with the modern web development concepts promoted by FastAPI. So using Uvicorn as a FastAPI web server is a good choice.
https://avoid.overfit.cn/post/c23627999f1a44689000105eb315180e
Author: auliyafirdaus