ChengZi
2 years ago
1 changed files with
11 additions and
8 deletions
-
timm_image.py
|
|
@ -31,7 +31,7 @@ from torch import nn |
|
|
|
from PIL import Image as PILImage |
|
|
|
|
|
|
|
import timm |
|
|
|
from timm.data import create_transform |
|
|
|
from timm.data import create_transform, resolve_data_config |
|
|
|
from timm.models import create_model, get_pretrained_cfg |
|
|
|
|
|
|
|
import warnings |
|
|
@ -90,6 +90,7 @@ class TimmImage(NNOperator): |
|
|
|
device=self.device, |
|
|
|
num_classes=num_classes |
|
|
|
) |
|
|
|
try: |
|
|
|
self.tfms = create_transform( |
|
|
|
input_size=self.config['input_size'], |
|
|
|
interpolation=self.config['interpolation'], |
|
|
@ -97,6 +98,8 @@ class TimmImage(NNOperator): |
|
|
|
std=self.config['std'], |
|
|
|
crop_pct=self.config['crop_pct'] |
|
|
|
) |
|
|
|
except: |
|
|
|
self.tfms = create_transform(**resolve_data_config({}, model=self.model)) |
|
|
|
self.skip_tfms = skip_preprocess |
|
|
|
else: |
|
|
|
log.warning('The operator is initialized without specified model.') |
|
|
|