Skip to content

Easy & Modular Computer Vision Detectors, Trackers & SAM - Run YOLOv9,v8,v7,v6,v5,R,X in under 10 lines of code.

License

Notifications You must be signed in to change notification settings

augmentedstartups/AS-One

Folders and files

NameName
Last commit message
Last commit date

Latest commit

 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 

Repository files navigation

AS-One : A Modular Library for YOLO Object Detection and Object Tracking

start with why GPLv3 License

Table of Contents

  1. Introduction
  2. Prerequisites
  3. Clone the Repo
  4. Installation
  5. Running AS-One
  6. Sample Code Snippets
  7. Model Zoo

1. Introduction

==UPDATE: YOLO-NAS is OUT==

AS-One is a python wrapper for multiple detection and tracking algorithms all at one place. Different trackers such as ByteTrack, DeepSORT or NorFair can be integrated with different versions of YOLO with minimum lines of code. This python wrapper provides YOLO models in ONNX, PyTorch & CoreML flavors. We plan to offer support for future versions of YOLO when they get released.

This is One Library for most of your computer vision needs.

If you would like to dive deeper into YOLO Object Detection and Tracking, then check out our courses and projects

Watch the step-by-step tutorial

2. Prerequisites

3. Clone the Repo

Navigate to an empty folder of your choice.

git clone https://github.com/augmentedstartups/AS-One.git

Change Directory to AS-One

cd AS-One

4. Installation

For Linux
python3 -m venv .env
source .env/bin/activate

pip install -r requirements.txt

# for CPU
pip install torch torchvision
# for GPU
pip install torch torchvision --extra-index-url https://download.pytorch.org/whl/cu113
For Windows 10/11
python -m venv .env
.env\Scripts\activate
pip install numpy Cython
pip install lap
pip install -e git+https://github.com/samson-wang/cython_bbox.git#egg=cython-bbox

pip install asone onnxruntime-gpu==1.12.1
pip install typing_extensions==4.7.1
pip install super-gradients==3.1.3
# for CPU
pip install torch torchvision

# for GPU
pip install torch torchvision --extra-index-url https://download.pytorch.org/whl/cu113
or
pip install torch==1.10.1+cu113 torchvision==0.11.2+cu113 torchaudio===0.10.1+cu113 -f https://download.pytorch.org/whl/cu113/torch_stable.html
For MacOS
python3 -m venv .env
source .env/bin/activate


pip install -r requirements.txt

# for CPU
pip install torch torchvision

5. Running AS-One

Run main.py to test tracker on data/sample_videos/test.mp4 video

python main.py data/sample_videos/test.mp4

Run in Google Colab

Open In Colab

6. Sample Code Snippets

6.1. Object Detection
from asone import ASOne, VideoReader, YOLOV7_PYTORCH

detector = ASOne(detector=YOLOV7_PYTORCH, use_cuda=True) # Set use_cuda to False for cpu
cap = VideoReader('data/sample_videos/test.mp4')

for frame in cap:
    dets, img_info = detector.detect(frame)
    frame = ASOne.draw_bboxes(frame, dets)

Run the asone/demo_detector.py to test detector.

# run on gpu
python -m asone.demo_detector data/sample_videos/test.mp4

# run on cpu
python -m asone.demo_detector data/sample_videos/test.mp4 --cpu
6.1.1 Use Custom Trained Weights for Detector

Use your custom weights of a detector model trained on custom data by simply providing path of the weights file.

from asone import ASOne, VideoReader, YOLOV7_PYTORCH

detector = ASOne(detector=YOLOV7_PYTORCH, weights='data/custom_weights/yolov7_custom.pt', use_cuda=True) # Set use_cuda to False for cpu
cap = VideoReader('data/sample_videos/license_video.mp4')

for frame in cap:
    dets, img_info = detector.detect(frame)
    frame = ASOne.draw_bboxes(frame, dets, class_names=['license_plate'])
6.1.2. Changing Detector Models

Change detector by simply changing detector flag. The flags are provided in benchmark tables.

  • Our library now supports YOLOv5, YOLOv7, and YOLOv8 on macOS.
# Change detector
detector = ASOne(detector=asone.YOLOX_S_PYTORCH, use_cuda=True)

# For macOs
# YOLO5
detector = ASOne(detector=asone.YOLOV5X_MLMODEL)
# YOLO7
detector = ASOne(detector=asone.YOLOV7_MLMODEL)
# YOLO8
detector = ASOne(detector=asone.YOLOV8L_MLMODEL)
6.2. Object Tracking

Use tracker on sample video.

from asone import ASOne, BYTETRACK, YOLOV7_PYTORCH

# Instantiate Asone object
detect = ASOne(tracker=asone.BYTETRACK, detector=asone.YOLOV7_PYTORCH, use_cuda=True) #set use_cuda=False to use cpu
track = detect.track_video('data/sample_videos/test.mp4', filter_classes=['car'])

# Loop over track to retrieve outputs of each frame
for bbox_details, frame_details in track:
    frame, _ , _ = frame_details
    frame = ASOne.draw_bboxes(frame, bbox_details)
    # Do anything with bboxes here

[Note] Use can use custom weights for a detector model by simply providing path of the weights file. in ASOne class.

6.2.1 Changing Detector and Tracking Models

Change Tracker by simply changing the tracker flag.

The flags are provided in benchmark tables.

detect = ASOne(tracker=asone.BYTETRACK, detector=asone.YOLOV7_PYTORCH, use_cuda=True)
# Change tracker
detect = ASOne(tracker=asone.DEEPSORT, detector=asone.YOLOV7_PYTORCH, use_cuda=True)
# Change Detector
detect = ASOne(tracker=asone.DEEPSORT, detector=asone.YOLOX_S_PYTORCH, use_cuda=True)

Run the asone/demo_tracker.py to test detector.

# run on gpu
python -m asone.demo_tracker data/sample_videos/test.mp4

# run on cpu
python -m asone.demo_tracker data/sample_videos/test.mp4 --cpu
6.3. Text Detection

Sample code to detect text on an image

# Detect and recognize text
from asone import ASOne, utils, CRAFT, EASYOCR
import cv2

ocr = ASOne(detector=CRAFT, recognizer=EASYOCR, use_cuda=True) # Set use_cuda to False for cpu
img = cv2.imread('data/sample_imgs/sample_text.jpeg')
results = ocr.detect_text(img)
img = utils.draw_text(img, results)

Use Tracker on Text

from asone import ASOne, DEEPSORT, CRAFT, EASYOCR

# Instantiate Asone object
detect = ASOne(tracker=DEEPSORT, detector=CRAFT, recognizer=EASYOCR, use_cuda=True) #set use_cuda=False to use cpu
track = detect.track_video('data/sample_videos/GTA_5-Unique_License_Plate.mp4')

# Loop over track to retrieve outputs of each frame
for bbox_details, frame_details in track:
    frame, _, _ = frame_details
    frame = ASOne.draw_bboxes(frame, bbox_details)

    # Do anything with bboxes here

Run the asone/demo_ocr.py to test ocr.

# run on gpu
 python -m asone.demo_ocr data/sample_videos/GTA_5-Unique_License_Plate.mp4

# run on cpu
 python -m asone.demo_ocr data/sample_videos/GTA_5-Unique_License_Plate.mp4 --cpu
6.4. Pose Estimation

Sample code to estimate pose on an image

# Pose Estimation
from asone import PoseEstimator, YOLOV8M_POSE, utils
import cv2

pose_estimator = PoseEstimator(estimator_flag=YOLOV8M_POSE, use_cuda=True) #set use_cuda=False to use cpu
img = cv2.imread('data/sample_imgs/test2.jpg')
kpts = pose_estimator.estimate_image(img)
img = utils.draw_kpts(img, kpts)
  • Now you can use Yolov8 and Yolov7-w6 for pose estimation. The flags are provided in benchmark tables.
# Pose Estimation on video
from asone import PoseEstimator, YOLOV7_W6_POSE, utils

pose_estimator = PoseEstimator(estimator_flag=YOLOV7_W6_POSE, use_cuda=True) #set use_cuda=False to use cpu
estimator = pose_estimator.estimate_video('data/sample_videos/football1.mp4')
for kpts, frame_details in estimator:
    frame, _, __ = frame_details
    img = utils.draw_kpts(frame, kpts)
    # Do anything with kpts here

Run the asone/demo_pose_estimator.py to test Pose estimation.

# run on gpu
 python -m asone.demo_pose_estimator data/sample_videos/football1.mp4

# run on cpu
 python -m asone.demo_pose_estimator data/sample_videos/football1.mp4 --cpu
6.5. Segmentation
from asone import ASOne, YOLOV7_PYTORCH, SAM

detect = ASOne(detector=YOLOV7_PYTORCH, segmentor=SAM, use_cuda=True) #set use_cuda=False to use cpu
track = detect.detect_video('data/sample_videos/test.mp4', filter_classes=['car'])

for bbox_details, frame_details in track:
    frame, frame_num, fps = frame_details
    frame = ASOne.draw_masks(frame, bbox_details) # Draw masks

To setup ASOne using Docker follow instructions given in docker setup

ToDo

  • First Release
  • Import trained models
  • Simplify code even further
  • Updated for YOLOv8
  • OCR and Counting
  • OCSORT, StrongSORT, MoTPy
  • M1/2 Apple Silicon Compatibility
  • Pose Estimation YOLOv7/v8
  • YOLO-NAS
  • SAM Integration
Offered By: Maintained By:
AugmentedStarups AxcelerateAI