Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 5 additions & 5 deletions rtmlib/tools/object_detection/yolox.py
Original file line number Diff line number Diff line change
Expand Up @@ -28,7 +28,7 @@ class YOLOX(BaseTool):
def __init__(self,
onnx_model: str,
model_input_size: tuple = (640, 640),
mode: str = 'human',
det_mode: str = 'human',
nms_thr=0.45,
score_thr=0.7,
backend: str = 'onnxruntime',
Expand All @@ -37,7 +37,7 @@ def __init__(self,
model_input_size,
backend=backend,
device=device)
self.mode = mode
self.det_mode = det_mode
self.nms_thr = nms_thr
self.score_thr = score_thr

Expand Down Expand Up @@ -157,11 +157,11 @@ def postprocess(
isbbox = [i for i in isscore]
final_boxes = final_boxes[isbbox]

if self.mode == 'multiclass':
if self.det_mode == 'multiclass':
return final_boxes, final_cls_inds
elif self.mode == 'human':
elif self.det_mode == 'human':
return final_boxes
else:
raise NotImplementedError(
f'Mode must be \'human\' or \'multiclass\': {self.mode} is not supported.'
f'det_mode must be \'human\' or \'multiclass\': {self.det_mode} is not supported.'
)
2 changes: 1 addition & 1 deletion rtmlib/tools/solution/animal.py
Original file line number Diff line number Diff line change
Expand Up @@ -89,7 +89,7 @@ def __init__(self,
det_input_size = self.MODE[mode]['det_input_size']

self.det_model = YOLOX(det,
mode='multiclass',
det_mode='multiclass',
model_input_size=det_input_size,
backend=backend,
device=device)
Expand Down
2 changes: 1 addition & 1 deletion rtmlib/tools/solution/custom.py
Original file line number Diff line number Diff line change
Expand Up @@ -120,7 +120,7 @@ def __init__(self,
det_class = getattr(rtmlib_module, det_class)
self.det_model = det_class(det,
model_input_size=det_input_size,
mode=det_mode,
det_mode=det_mode,
backend=backend,
device=device)
self.det_mode = det_mode
Expand Down