Skip to content

Commit

Permalink
Add demo option
Browse files Browse the repository at this point in the history
  • Loading branch information
@PINTO0309
PINTO0309 committed Nov 13, 2024
1 parent 9c36a8a commit 4d60e0d
Show file tree
Hide file tree
Showing 2 changed files with 73 additions and 12 deletions.
14 changes: 13 additions & 1 deletion 459_YOLOv9-Wholebody25/README.md
View file
Original file line number Diff line number Diff line change
Expand Up @@ -127,11 +127,15 @@ The use of [CD-COCO: Complex Distorted COCO database for Scene-Context-Aware com
[-it] \
[-dvw] \
[-dwk] \
[-ost] \
[-ast] \
[-dnm] \
[-dgm] \
[-dlr] \
[-dhm] \
[-oyt]
[-drc] \
[-oyt] \
[-bblw]

options:
-h, --help
Expand All @@ -154,6 +158,10 @@ The use of [CD-COCO: Complex Distorted COCO database for Scene-Context-Aware com
-dwk, --disable_waitKey
Disable cv2.waitKey(). When you want to process a batch of still images,
disable key-input wait and process them continuously.
-ost OBJECT_SOCRE_THRESHOLD, --object_socre_threshold OBJECT_SOCRE_THRESHOLD
The detection score threshold for object detection. Default: 0.35
-ast ATTRIBUTE_SOCRE_THRESHOLD, --attribute_socre_threshold ATTRIBUTE_SOCRE_THRESHOLD
The attribute score threshold for object detection. Default: 0.70
-dnm, --disable_generation_identification_mode
Disable generation identification mode.
(Press N on the keyboard to switch modes)
Expand All @@ -166,8 +174,12 @@ The use of [CD-COCO: Complex Distorted COCO database for Scene-Context-Aware com
-dhm, --disable_headpose_identification_mode
Disable HeadPose identification mode.
(Press P on the keyboard to switch modes)
-drc [DISABLE_RENDER_CLASSIDS ...], --disable_render_classids [DISABLE_RENDER_CLASSIDS ...]
Class ID to disable bounding box drawing. List[int]. e.g. -drc 17 18 19
-oyt, --output_yolo_format_text
Output YOLO format texts and images.
-bblw BOUNDING_BOX_LINE_WIDTH, --bounding_box_line_width BOUNDING_BOX_LINE_WIDTH
Bounding box line width. Default: 2
```

- YOLOv9-Wholebody25 - N - Swish/SiLU (PINTO original implementation, 2.4 MB)
Expand Down
71 changes: 60 additions & 11 deletions 459_YOLOv9-Wholebody25/demo/demo_yolov9_onnx_wholebody25.py
100644 → 100755
View file
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,7 @@
from enum import Enum
from pathlib import Path
from dataclasses import dataclass
from argparse import ArgumentParser
from argparse import ArgumentParser, ArgumentTypeError
from typing import Tuple, Optional, List, Dict
import importlib.util
from abc import ABC, abstractmethod
Expand Down Expand Up @@ -670,6 +670,13 @@ def draw_dashed_rectangle(

def main():
parser = ArgumentParser()

def check_positive(value):
ivalue = int(value)
if ivalue < 2:
raise ArgumentTypeError(f"Invalid Value: {ivalue}. Please specify an integer of 2 or greater.")
return ivalue

parser.add_argument(
'-m',
'--model',
Expand Down Expand Up @@ -724,6 +731,22 @@ def main():
'When you want to process a batch of still images, '+
' disable key-input wait and process them continuously.',
)
parser.add_argument(
'-ost',
'--object_socre_threshold',
type=float,
default=0.35,
help=\
'The detection score threshold for object detection. Default: 0.35',
)
parser.add_argument(
'-ast',
'--attribute_socre_threshold',
type=float,
default=0.75,
help=\
'The attribute score threshold for object detection. Default: 0.70',
)
parser.add_argument(
'-dnm',
'--disable_generation_identification_mode',
Expand Down Expand Up @@ -752,13 +775,30 @@ def main():
help=\
'Disable HeadPose identification mode. (Press P on the keyboard to switch modes)',
)
parser.add_argument(
'-drc',
'--disable_render_classids',
type=int,
nargs="*",
default=[],
help=\
'Class ID to disable bounding box drawing. List[int]. e.g. -drc 17 18 19',
)
parser.add_argument(
'-oyt',
'--output_yolo_format_text',
action='store_true',
help=\
'Output YOLO format texts and images.',
)
parser.add_argument(
'-bblw',
'--bounding_box_line_width',
type=check_positive,
default=2,
help=\
'Bounding box line width. Default: 2',
)
args = parser.parse_args()

# runtime check
Expand All @@ -784,14 +824,18 @@ def main():
video: str = args.video
images_dir: str = args.images_dir
disable_waitKey: bool = args.disable_waitKey
object_socre_threshold: float = args.object_socre_threshold
attribute_socre_threshold: float = args.attribute_socre_threshold
disable_generation_identification_mode: bool = args.disable_generation_identification_mode
disable_gender_identification_mode: bool = args.disable_gender_identification_mode
disable_left_and_right_hand_identification_mode: bool = args.disable_left_and_right_hand_identification_mode
disable_headpose_identification_mode: bool = args.disable_headpose_identification_mode
disable_render_classids: List[int] = args.disable_render_classids
output_yolo_format_text: bool = args.output_yolo_format_text
execution_provider: str = args.execution_provider
inference_type: str = args.inference_type
inference_type = inference_type.lower()
bounding_box_line_width: int = args.bounding_box_line_width
providers: List[Tuple[str, Dict] | str] = None

if execution_provider == 'cpu':
Expand Down Expand Up @@ -842,8 +886,8 @@ def main():
model = YOLOv9(
runtime=runtime,
model_path=model_file,
obj_class_score_th=0.35,
attr_class_score_th=0.70,
obj_class_score_th=object_socre_threshold,
attr_class_score_th=attribute_socre_threshold,
providers=providers,
)

Expand Down Expand Up @@ -871,6 +915,8 @@ def main():

file_paths_count = -1
movie_frame_count = 0
white_line_width = bounding_box_line_width
colored_line_width = white_line_width - 1
while True:
image: np.ndarray = None
if file_paths is not None:
Expand Down Expand Up @@ -908,6 +954,9 @@ def main():
classid: int = box.classid
color = (255,255,255)

if classid in disable_render_classids:
continue

if classid == 0:
# Body
if not disable_gender_identification_mode:
Expand Down Expand Up @@ -985,8 +1034,8 @@ def main():
dash_length=10
)
else:
cv2.rectangle(debug_image, (box.x1, box.y1), (box.x2, box.y2), (255,255,255), 3)
cv2.rectangle(debug_image, (box.x1, box.y1), (box.x2, box.y2), color, 2)
cv2.rectangle(debug_image, (box.x1, box.y1), (box.x2, box.y2), (255,255,255), white_line_width)
cv2.rectangle(debug_image, (box.x1, box.y1), (box.x2, box.y2), color, colored_line_width)

elif classid == 7:
if box.head_pose == -1:
Expand All @@ -999,8 +1048,8 @@ def main():
dash_length=10
)
else:
cv2.rectangle(debug_image, (box.x1, box.y1), (box.x2, box.y2), (255,255,255), 3)
cv2.rectangle(debug_image, (box.x1, box.y1), (box.x2, box.y2), color, 2)
cv2.rectangle(debug_image, (box.x1, box.y1), (box.x2, box.y2), (255,255,255), white_line_width)
cv2.rectangle(debug_image, (box.x1, box.y1), (box.x2, box.y2), color, colored_line_width)

elif classid == 21:
if box.handedness == -1:
Expand All @@ -1013,12 +1062,12 @@ def main():
dash_length=10
)
else:
cv2.rectangle(debug_image, (box.x1, box.y1), (box.x2, box.y2), (255,255,255), 3)
cv2.rectangle(debug_image, (box.x1, box.y1), (box.x2, box.y2), color, 2)
cv2.rectangle(debug_image, (box.x1, box.y1), (box.x2, box.y2), (255,255,255), white_line_width)
cv2.rectangle(debug_image, (box.x1, box.y1), (box.x2, box.y2), color, colored_line_width)

else:
cv2.rectangle(debug_image, (box.x1, box.y1), (box.x2, box.y2), (255,255,255), 3)
cv2.rectangle(debug_image, (box.x1, box.y1), (box.x2, box.y2), color, 2)
cv2.rectangle(debug_image, (box.x1, box.y1), (box.x2, box.y2), (255,255,255), white_line_width)
cv2.rectangle(debug_image, (box.x1, box.y1), (box.x2, box.y2), color, colored_line_width)

# Attributes text
generation_txt = ''
Expand Down

0 comments on commit 4d60e0d

Please sign in to comment.