logo
Browse Source

modify3

Signed-off-by: xujinling <jinling.xu@zilliz.com>
main
xujinling 2 years ago
parent
commit
dcba73be8f
  1. 2
      video_swin_transformer.py

2
video_swin_transformer.py

@ -10,6 +10,7 @@ from towhee.operator.base import NNOperator
from towhee.types.video_frame import VideoFrame from towhee.types.video_frame import VideoFrame
from towhee.models.utils.video_transforms import transform_video, get_configs from towhee.models.utils.video_transforms import transform_video, get_configs
from towhee.models.video_swin_transformer import video_swin_transformer from towhee.models.video_swin_transformer import video_swin_transformer
from towhee.models.video_swin_transformer.get_configs import configs
log = logging.getLogger() log = logging.getLogger()
@ -41,6 +42,7 @@ class VideoSwinTransformer(NNOperator):
self.model_name = model_name self.model_name = model_name
self.skip_preprocess = skip_preprocess self.skip_preprocess = skip_preprocess
self.topk = topk self.topk = topk
self.model_configs = configs(self.model_name)
if classmap is None: if classmap is None:
class_file = os.path.join(str(Path(__file__).parent), self.model_configs['labels_file_name']) class_file = os.path.join(str(Path(__file__).parent), self.model_configs['labels_file_name'])
with open(class_file, 'r') as f: with open(class_file, 'r') as f:

Loading…
Cancel
Save