Thastp commited on
Commit
8d76e0c
·
verified ·
1 Parent(s): a87fd2f

Update image_processing_efficientnet.py

Browse files
Files changed (1) hide show
  1. image_processing_efficientnet.py +26 -26
image_processing_efficientnet.py CHANGED
@@ -1,27 +1,27 @@
1
- from transformers.image_processing_utils import BaseImageProcessor, BatchFeature
2
- from configuration_efficientnet import MODEL_NAMES
3
- from timm import create_model
4
- from timm.data import resolve_data_config
5
- from timm.data.transforms_factory import create_transform
6
-
7
- class EfficientNetImageProcessor(BaseImageProcessor):
8
- model_input_names = ["pixel_values"]
9
-
10
- def __init__(self,
11
- model_name: str,
12
- **kwargs
13
- ):
14
- super().__init__(**kwargs)
15
-
16
- self.model_name = model_name
17
- self.config = resolve_data_config({}, model=create_model(model_name, pretrained=False))
18
-
19
-
20
- def preprocess(self, image):
21
- transforms = create_transform(**self.config)
22
- data = {'pixel_values': transforms(image).unsqueeze(0)}
23
- return BatchFeature(data=data)
24
-
25
- __all__ = [
26
- "EfficientNetImageProcessor"
27
  ]
 
1
+ from transformers.image_processing_utils import BaseImageProcessor, BatchFeature
2
+ from .configuration_efficientnet import MODEL_NAMES
3
+ from timm import create_model
4
+ from timm.data import resolve_data_config
5
+ from timm.data.transforms_factory import create_transform
6
+
7
+ class EfficientNetImageProcessor(BaseImageProcessor):
8
+ model_input_names = ["pixel_values"]
9
+
10
+ def __init__(self,
11
+ model_name: str,
12
+ **kwargs
13
+ ):
14
+ super().__init__(**kwargs)
15
+
16
+ self.model_name = model_name
17
+ self.config = resolve_data_config({}, model=create_model(model_name, pretrained=False))
18
+
19
+
20
+ def preprocess(self, image):
21
+ transforms = create_transform(**self.config)
22
+ data = {'pixel_values': transforms(image).unsqueeze(0)}
23
+ return BatchFeature(data=data)
24
+
25
+ __all__ = [
26
+ "EfficientNetImageProcessor"
27
  ]