diff --git a/modelscope/models/cv/realtime_object_detection/realtime_detector.py b/modelscope/models/cv/realtime_object_detection/realtime_detector.py index b147f769..2b4b3f8c 100644 --- a/modelscope/models/cv/realtime_object_detection/realtime_detector.py +++ b/modelscope/models/cv/realtime_object_detection/realtime_detector.py @@ -1,3 +1,4 @@ +# Copyright (c) Alibaba, Inc. and its affiliates. import argparse import logging as logger import os @@ -48,6 +49,7 @@ class RealtimeDetector(TorchModel): self.nmsthre = self.exp.nmsthre self.test_size = self.exp.test_size self.preproc = ValTransform(legacy=False) + self.label_mapping = self.config['labels'] def inference(self, img): with torch.no_grad(): @@ -81,5 +83,8 @@ class RealtimeDetector(TorchModel): bboxes = outputs[0][:, 0:4].cpu().numpy() / self.ratio scores = outputs[0][:, 5].cpu().numpy() labels = outputs[0][:, 6].cpu().int().numpy() + pred_label_names = [] + for lab in labels: + pred_label_names.append(self.label_mapping[lab]) - return bboxes, scores, labels + return bboxes, scores, pred_label_names diff --git a/modelscope/models/cv/realtime_object_detection/yolox/exp/default/yolox_nano.py b/modelscope/models/cv/realtime_object_detection/yolox/exp/default/yolox_nano.py index 330eef16..7bada485 100644 --- a/modelscope/models/cv/realtime_object_detection/yolox/exp/default/yolox_nano.py +++ b/modelscope/models/cv/realtime_object_detection/yolox/exp/default/yolox_nano.py @@ -42,5 +42,6 @@ class YoloXNanoExp(YoloXExp): act=self.act, depthwise=True) self.model = YOLOX(backbone, head) - + self.model.apply(init_yolo) + self.model.head.initialize_biases(1e-2) return self.model diff --git a/modelscope/pipelines/cv/realtime_object_detection_pipeline.py b/modelscope/pipelines/cv/realtime_object_detection_pipeline.py index 629720d1..9f558f88 100644 --- a/modelscope/pipelines/cv/realtime_object_detection_pipeline.py +++ b/modelscope/pipelines/cv/realtime_object_detection_pipeline.py @@ -1,3 +1,4 @@ +# Copyright (c) Alibaba, Inc. and its affiliates. import os.path as osp from typing import Any, Dict, List, Union