diff --git a/README.md b/README.md index ba7c7ea..a7002fb 100644 --- a/README.md +++ b/README.md @@ -61,6 +61,8 @@ pip install cython-bbox asone onnxruntime-gpu==1.12.1 pip install super-gradients==3.1.1 # for CPU pip install torch torchvision +# for GPU +pip install torch torchvision --extra-index-url https://download.pytorch.org/whl/cu113 ``` @@ -70,7 +72,8 @@ pip install torch torchvision ```shell python -m venv .env .env\Scripts\activate -pip install numpy Cython +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 @@ -96,10 +99,6 @@ pip install cython-bbox asone pip install super-gradients==3.1.1 # for CPU pip install torch torchvision - -# for GPU -pip install torch torchvision --extra-index-url https://download.pytorch.org/whl/cu113 - ``` @@ -213,11 +212,11 @@ detector = ASOne(detector=asone.YOLOX_S_PYTORCH, use_cuda=True) # For macOs # YOLO5 -detector = ASOne(detector=asone.YOLOV5X_MLMODEL, use_cuda=True) +detector = ASOne(detector=asone.YOLOV5X_MLMODEL) # YOLO7 -detector = ASOne(detector=asone.YOLO7_MLMODEL, use_cuda=True) +detector = ASOne(detector=asone.YOLOV7_MLMODEL) # YOLO8 -detector = ASOne(detector=asone.YOLOV8L_MLMODEL, use_cuda=True) +detector = ASOne(detector=asone.YOLOV8L_MLMODEL) ``` @@ -425,7 +424,6 @@ To setup ASOne using Docker follow instructions given in [docker setup](asone/li - [x] YOLO-NAS - [ ] SAM Integration - |Offered By: |Maintained By:| |-------------|-------------| |[![AugmentedStarups](https://user-images.githubusercontent.com/107035454/195115263-d3271ef3-973b-40a4-83c8-0ade8727dd40.png)](https://augmentedstartups.com)|[![AxcelerateAI](https://user-images.githubusercontent.com/107035454/195114870-691c8a52-fcf0-462e-9e02-a720fc83b93f.png)](https://axcelerate.ai/)| diff --git a/asone/demo_detector.py b/asone/demo_detector.py index 0a5c4fb..b39c19b 100644 --- a/asone/demo_detector.py +++ b/asone/demo_detector.py @@ -5,6 +5,9 @@ import argparse import time import os +import sys +import torch + def main(args): filter_classes = args.filter_classes @@ -15,8 +18,17 @@ def main(args): if filter_classes: filter_classes = filter_classes.split(',') + if args.use_cuda and torch.cuda.is_available(): + args.use_cuda = True + else: + args.use_cuda = False - detector = ASOne(asone.YOLOV7_PYTORCH, weights=args.weights, use_cuda=args.use_cuda) + if sys.platform.startswith('darwin'): + detector = asone.YOLOV7_MLMODEL + else: + detector = asone.YOLOV7_PYTORCH + + detector = ASOne(detector, weights=args.weights, use_cuda=args.use_cuda) cap = cv2.VideoCapture(video_path) width = cap.get(cv2.CAP_PROP_FRAME_WIDTH) diff --git a/asone/demo_tracker.py b/asone/demo_tracker.py index 733a5bc..f24ef6e 100644 --- a/asone/demo_tracker.py +++ b/asone/demo_tracker.py @@ -18,7 +18,7 @@ def main(args): detect = ASOne(tracker=asone.BYTETRACK, detector=asone.YOLOV7_PYTORCH, - use_cuda=True) + use_cuda=args.use_cuda) track = detect.track_video(video_path, output_dir=args.output_path, save_result=args.save, display=args.display, diff --git a/asone/pose_estimators/yolov7_pose/models/common.py b/asone/pose_estimators/yolov7_pose/models/common.py index 6f28cd5..532932d 100644 --- a/asone/pose_estimators/yolov7_pose/models/common.py +++ b/asone/pose_estimators/yolov7_pose/models/common.py @@ -12,10 +12,10 @@ from PIL import Image from torch.cuda import amp -from utils.datasets import letterbox -from utils.general import non_max_suppression, make_divisible, scale_coords, increment_path, xyxy2xywh -from utils.plots import color_list, plot_one_box -from utils.torch_utils import time_synchronized +from asone.pose_estimators.yolov7_pose.utils.datasets import letterbox +from asone.pose_estimators.yolov7_pose.utils.general import non_max_suppression, make_divisible, scale_coords, increment_path, xyxy2xywh +from asone.pose_estimators.yolov7_pose.utils.plots import color_list, plot_one_box +from asone.pose_estimators.yolov7_pose.utils.torch_utils import time_synchronized ##### basic #### diff --git a/asone/pose_estimators/yolov7_pose/models/experimental.py b/asone/pose_estimators/yolov7_pose/models/experimental.py index 46f2ba6..90e126b 100644 --- a/asone/pose_estimators/yolov7_pose/models/experimental.py +++ b/asone/pose_estimators/yolov7_pose/models/experimental.py @@ -3,9 +3,9 @@ import torch import torch.nn as nn -from models.common import Conv, DWConv -from utils.google_utils import attempt_download - +from .common import Conv, DWConv +from asone.pose_estimators.yolov7_pose.utils.google_utils import attempt_download +from asone.pose_estimators.yolov7_pose.utils.yolov7_pose_utils import yolov7_in_syspath class CrossConv(nn.Module): # Cross Convolution Downsample @@ -239,7 +239,8 @@ def attempt_load(weights, map_location=None): model = Ensemble() for w in weights if isinstance(weights, list) else [weights]: attempt_download(w) - ckpt = torch.load(w, map_location=map_location) # load + with yolov7_in_syspath(): + ckpt = torch.load(w, map_location=map_location) # load model.append(ckpt['ema' if ckpt.get('ema') else 'model'].float().fuse().eval()) # FP32 model # Compatibility updates diff --git a/asone/pose_estimators/yolov7_pose/utils/datasets.py b/asone/pose_estimators/yolov7_pose/utils/datasets.py index 094fa41..0107f60 100644 --- a/asone/pose_estimators/yolov7_pose/utils/datasets.py +++ b/asone/pose_estimators/yolov7_pose/utils/datasets.py @@ -28,7 +28,7 @@ from .general import check_requirements, xyxy2xywh, xywh2xyxy, xywhn2xyxy, xyn2xy, segment2box, segments2boxes, \ resample_segments, clean_str -from utils.torch_utils import torch_distributed_zero_first +from .torch_utils import torch_distributed_zero_first # Parameters help_url = 'https://github.com/ultralytics/yolov5/wiki/Train-Custom-Data' diff --git a/asone/pose_estimators/yolov7_pose/utils/yolov7_pose_utils.py b/asone/pose_estimators/yolov7_pose/utils/yolov7_pose_utils.py new file mode 100644 index 0000000..7ef27bb --- /dev/null +++ b/asone/pose_estimators/yolov7_pose/utils/yolov7_pose_utils.py @@ -0,0 +1,20 @@ +import contextlib +from pathlib import Path +import sys + +@contextlib.contextmanager +def yolov7_in_syspath(): + """ + Temporarily add yolov5 folder to `sys.path`. + + torch.hub handles it in the same way: https://github.com/pytorch/pytorch/blob/75024e228ca441290b6a1c2e564300ad507d7af6/torch/hub.py#L387 + + Proper fix for: #22, #134, #353, #1155, #1389, #1680, #2531, #3071 + No need for such workarounds: #869, #1052, #2949 + """ + yolov7_folder_dir = str(Path(__file__).parents[1].absolute()) + try: + sys.path.insert(0, yolov7_folder_dir) + yield + finally: + sys.path.remove(yolov7_folder_dir) \ No newline at end of file diff --git a/asone/pose_estimators/yolov7_pose/yolov7.py b/asone/pose_estimators/yolov7_pose/yolov7.py index 01988c3..5ec29ce 100644 --- a/asone/pose_estimators/yolov7_pose/yolov7.py +++ b/asone/pose_estimators/yolov7_pose/yolov7.py @@ -5,19 +5,13 @@ import matplotlib.pyplot as plt from torchvision import transforms from .utils.datasets import letterbox -from utils.torch_utils import select_device -from models.experimental import attempt_load +from .utils.torch_utils import select_device +from .models.experimental import attempt_load from .utils.general import non_max_suppression_kpt,strip_optimizer,xyxy2xywh from .utils.plots import output_to_keypoint, plot_skeleton_kpts,colors,plot_one_box_kpt +import os +import sys -def scale_bboxes(bboxes, org_img_shape, resized_img_shape): - # Rescaling Bounding Boxes - # print(bboxes[:, :2]) - # exit() - bboxes[:, :2] /= np.array([resized_img_shape[1], resized_img_shape[0]]) - bboxes[:, :2] *= np.array([org_img_shape[1], org_img_shape[0]]) - - return bboxes class Yolov7PoseEstimator: def __init__(self, weights="yolov7-w6-pose.pt", use_cuda=True): diff --git a/main.py b/main.py index 1788cc6..0d8d665 100644 --- a/main.py +++ b/main.py @@ -1,16 +1,29 @@ +import sys import argparse import asone from asone import ASOne +import torch + def main(args): filter_classes = args.filter_classes if filter_classes: filter_classes = ['person'] - + # Check if cuda available + if args.use_cuda and torch.cuda.is_available(): + args.use_cuda = True + else: + args.use_cuda = False + + if sys.platform.startswith('darwin'): + detector = asone.YOLOV7_MLMODEL + else: + detector = asone.YOLOV7_PYTORCH + detect = ASOne( tracker=asone.BYTETRACK, - detector=asone.YOLOV7_PYTORCH, + detector=detector, weights=args.weights, use_cuda=args.use_cuda ) diff --git a/requirements.txt b/requirements.txt index f72aaba..954897f 100644 --- a/requirements.txt +++ b/requirements.txt @@ -12,14 +12,10 @@ pandas tabulate typing-extensions==3.10.0.2 wheel -torch==1.12.1 -torchvision numpy==1.23.3 -Cython -ultralytics asone-ocr motpy - +ultralytics==8.0.109 torchreid==0.2.5 tensorboard protobuf==3.20 diff --git a/setup.py b/setup.py index 3578d3d..9875f7f 100644 --- a/setup.py +++ b/setup.py @@ -9,7 +9,7 @@ URL = 'https://github.com/axcelerateai/asone' DOWNLOAD_URL = URL -VERSION = '0.3.0' +VERSION = '0.3.2' with open('README.md') as f: long_description = f.read()