towhee
/
efficientnet-image-embedding
copied
4 changed files with 26 additions and 83 deletions
@ -1,22 +0,0 @@ |
|||
# Copyright 2021 Zilliz. All rights reserved. |
|||
# |
|||
# Licensed under the Apache License, Version 2.0 (the "License"); |
|||
# you may not use this file except in compliance with the License. |
|||
# You may obtain a copy of the License at |
|||
# |
|||
# http://www.apache.org/licenses/LICENSE-2.0 |
|||
# |
|||
# Unless required by applicable law or agreed to in writing, software |
|||
# distributed under the License is distributed on an "AS IS" BASIS, |
|||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
|||
# See the License for the specific language governing permissions and |
|||
# limitations under the License. |
|||
|
|||
|
|||
import os |
|||
|
|||
# For requirements. |
|||
try: |
|||
import efficientnet_pytorch |
|||
except ModuleNotFoundError: |
|||
os.system('pip install efficientnet_pytorch') |
@ -1,39 +0,0 @@ |
|||
# Copyright 2021 Zilliz. All rights reserved. |
|||
# |
|||
# Licensed under the Apache License, Version 2.0 (the "License"); |
|||
# you may not use this file except in compliance with the License. |
|||
# You may obtain a copy of the License at |
|||
# |
|||
# http://www.apache.org/licenses/LICENSE-2.0 |
|||
# |
|||
# Unless required by applicable law or agreed to in writing, software |
|||
# distributed under the License is distributed on an "AS IS" BASIS, |
|||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
|||
# See the License for the specific language governing permissions and |
|||
# limitations under the License. |
|||
|
|||
|
|||
import torch |
|||
import timm |
|||
|
|||
|
|||
class Model(): |
|||
""" |
|||
PyTorch model class |
|||
""" |
|||
def __init__(self, model_name: str, weights_path: str): |
|||
super().__init__() |
|||
if weights_path: |
|||
self._model = timm.create_model(model_name, checkpoint_path=weights_path, num_classes=0) |
|||
else: |
|||
self._model = timm.create_model(model_name, pretrained=True, num_classes=0) |
|||
self._model.eval() |
|||
|
|||
def __call__(self, img_tensor: torch.Tensor): |
|||
return self._model(img_tensor) |
|||
|
|||
def train(self): |
|||
""" |
|||
For training model |
|||
""" |
|||
pass |
Loading…
Reference in new issue