/yolor-trt

Fastapi + celery + yolor_trt

Primary LanguagePython

YOLOR-TRT

Empowering Vision with High-Speed Object Detection

Built with the tools and technologies:

Redis RabbitMQ Docker Python


Table of Contents

Overview

Yolor-trt is a pioneering open-source project providing an efficient, predictive personal protective equipment (PPE) detection system. This revolutionary technology leverages high-performance object detection models paired with convenient API endpoints, orchestrated services and asynchronous tasks handling. Designed for businesses, developers, and researchers, it ensures a safer workplace by accurately identifying PPE in visual content, delivering streamlined image processing, and facilitating rapid deployment.


Features

Feature Summary
โš™๏ธ Architecture
  • Based on the combination of FastAPI, TensorRT and Celery for distributed PPE detection using image data sent via API calls.
  • Task queue and message brokering are managed by RabbitMQ and Redis, ensuring a robust, high-performance system.
  • TensorRT environment is prepared through a Dockerfile, ensuring smooth operation of the system.
๐Ÿ”ฉ Code Quality
  • Consistent use of Python's PEP-8 conventions ensures code readability.
  • The code is clearly segmented into utilities, services, routes and controllers facilitating maintainability.
  • Use of Docker and docker-compose ensures consistent operation across different environments.
๐Ÿ“„ Documentation
  • The primary language used in the project is Python, with the presence of yml, py, txt, and names file types.
  • Dependencies are managed by pip and can be found in requirements.txt.
  • Containers are defined using Docker and orchestrated using docker-compose.yml.
๐Ÿ”Œ Integrations
  • Integrated with TensorRT for efficient implementation of the deep learning model.
  • Utilizes FastAPI for routing and API endpoints management.
  • Deploys RabbitMQ and Redis for task queuing and message brokering respectively.
๐Ÿงฉ Modularity
  • The project is modular with distinct directories for routes, services, controllers and utility functions.
  • Routes are managed by FastAPI routers, enhancing traceability and organization within the larger architectural framework.
  • Utility functions are abstracted out in model_utils.py for use across the application, encouraging code reusability.
๐Ÿงช Testing
  • The project utilizes PyTest for running tests, as indicated by the test commands in the documentation.
โšก๏ธ Performance
  • The project leverages NVIDIA's TensorRT models for efficient model inference, contributing to faster results.
  • The use of Non-Maximum Suppression (NMS) in model_utils.py helps in predicting accurate output by eliminating redundant bounding boxes.
๐Ÿ›ก๏ธ Security
  • Use of Docker containers, which provide added security through process isolation.
๐Ÿ“ฆ Dependencies
  • Managed by pip as defined in the requirements.txt file.
  • The system also relies on Docker, RabbitMQ, and Redis systems.

Project Structure

โ””โ”€โ”€ yolor-trt/
    โ”œโ”€โ”€ PPEDetection
    โ”‚   โ”œโ”€โ”€ .gitignore
    โ”‚   โ”œโ”€โ”€ Dockerfile
    โ”‚   โ”œโ”€โ”€ main.py
    โ”‚   โ”œโ”€โ”€ requirements.txt
    โ”‚   โ””โ”€โ”€ src
    โ”œโ”€โ”€ README.md
    โ””โ”€โ”€ docker-compose.yml

Project Index

YOLOR-TRT/
__root__
docker-compose.yml - The docker-compose.yml orchestrates service creation for a predictive personal protective equipment (PPE) detection system, encompassing a main PPE detection service, an associated worker service, and two additional required services: Redis and RabbitMQ
- These coordinated services ensure efficient message brokering and task queuing for robust, high-performance PPE detection.
PPEDetection
main.py - MaineCoon is a FastAPI-based application outlined in PPEDetection/main.py, offering a comprehensive routing solution within the project structure through the integration of the API router
- It ensures smooth navigation and efficient functionality of various endpoints, contributing significantly to the architectural integrity of the entire codebase.
requirements.txt - PPEDetection/requirements.txt lists the specific software dependencies required for the PPEDetection project
- It enables the setup of a uniform development environment and guarantees the smooth running of the project
- It includes packages necessary for computer vision tasks, web server operation, data serialization, and handling asynchronous tasks.
Dockerfile - The Dockerfile in the PPEDetection directory sets up an environment for running a TensorRT-based application
- It begins with an Nvidia TensorRT base image, installs necessary Python packages, copies project requirements, and installs them
- Additionally, it updates system packages and installs utilities required for multimedia processing.
src
routes
model_route.py - Model_route.py forms part of the routing system in the PPEDetection application
- It handles image data sent via API calls for inference, initiates processing, and manages responses
- It also retrieves processing results once ready, converting them to a viewable image format for return to the client.
routes.py - Routes.py in PPEDetection/src/routes acts as a connection point, integrating model_route with the APIRouter from FastAPI
- Predominantly, it assigns a specific "/model" prefix and "model" tag, enhancing the traceability and organization of the routes related to the model within the broader codebase architecture.
controllers
model_controller.py - Model Controller in the PPEDetection project facilitates image inference tasks
- Leveraging helper functions, it converts image bytes to string format, then dispatches this data for further processing
- Thus, it serves as an important touchpoint for managing image parsing and initiating subsequent AI analysis tasks.
queue
worker.py - Worker.py configures and initializes the Celery application within the Personal Protective Equipment (PPE) Detection project
- It leverages environment variables for the broker and backend URI setups
- Additionally, it includes task definitions from the tasks file located within the same queue module, facilitating distributed task execution across worker nodes.
tasks.py - The 'tasks.py' within the PPEDetection/src/queue directory serves as the primary component for managing machine learning model prediction tasks
- It utilizes an abstraction of the Celery's Task class to efficiently load models for the first task call, eliminating the need to reload for each subsequent task
- It also provides an inference function for image data processing.
utils
asserts
coco.names - The utility file 'coco.names' within the PPEDetection project serves as a predefined dictionary for object recognition
- It lists common objects from everyday items like 'bottle' or 'clock', to animals like 'dog', and transportation mediums like 'bicycle'
- It aids in identifying and labeling these objects within images or video frames for the personal protective equipment detection system.
helper
model_utils.py - The code in 'model_utils.py' provides various utility functions geared towards image processing and detection tasks
- Key functionalities include converting and resizing images, performing Non-Maximum Suppression (NMS) on prediction results, drawing bounding boxes for detection results, and handling image formats for further processing
- These functions support the primary tasks of the PPEDetection project.
services
trt_loader.py - The 'trt_loader.py' file in the PPEDetection project is responsible for managing NVIDIA's TensorRT models
- It prepares, executes, and retrieves results from the inference engine
- Specifically, it handles memory allocation for the engine's inputs and outputs, performs inference, and manages the transfer of data between host and device
- Furthermore, it also reorganizes the model's flattened outputs back to their original shapes.
model_service.py - The code in 'model_service.py' is responsible for coordinating the entire object detection process
- It manages the loading and configuration of the YOLOR model, preprocesses images for detection, handles detection and results from the model, and conducts post-processing activities including drawing bounding boxes and scaling results back to original image size.

Getting Started

Prerequisites

Before getting started with yolor-trt, ensure your runtime environment meets the following requirements:

  • Programming Language: Python
  • Package Manager: Pip
  • Container Runtime: Docker

Installation

Install yolor-trt using one of the following methods:

Build from source:

  1. Clone the yolor-trt repository:
โฏ git clone https://github.com/RTae/yolor-trt
  1. Navigate to the project directory:
โฏ cd yolor-trt
  1. Install the project dependencies:

Using pip  

โฏ pip install -r PPEDetection/requirements.txt

Using docker  

โฏ docker build -t RTae/yolor-trt .

Usage

Run yolor-trt using the following command: Using pip  

โฏ python {entrypoint}

Using docker  

โฏ docker run -it {image_name}