img2vec-pytorch - pypi Package Compare versions

Comparing version 1.0.0 to 1.0.1

img2vec_pytorch.egg-info/PKG-INFO

 Metadata-Version: 2.1
 Name: img2vec-pytorch
-Version: 1.0.0
+Version: 1.0.1
 Summary: Use pre-trained models in PyTorch to extract vector embeddings for any image
@@ -30,17 +30,25 @@ Home-page: https://github.com/christiansafka/img2vec
 ## Available models
- - Resnet-18 (CPU, GPU)
- 	- Returns vector length 512
- - Alexnet (CPU, GPU)
- 	- Returns vector length 4096
- - Vgg-11 (CPU, GPU)
-	- Returns vector length 4096
- - Densenet (CPU, GPU)
-	- Returns vector length 1024
+|Model name|Return vector length|
+|----|----|
+|Resnet-18|512|
+|Alexnet|4096|
+|Vgg-11|4096|
+|Densenet|1024|
+|efficientnet_b0|1280|
+|efficientnet_b1|1280|
+|efficientnet_b2|1408|
+|efficientnet_b3|1536|
+|efficientnet_b4|1792|
+|efficientnet_b5|2048|
+|efficientnet_b6|2304|
+|efficientnet_b7|2560|
 
 ## Installation
 
-Tested on Python 3.6
+Tested on Python 3.6 and torchvision 0.11.0 (nightly, 2021-09-25) 
 
 Requires Pytorch: http://pytorch.org/
 
+```conda install -c pytorch-nightly torchvision```
+
 ```pip install img2vec_pytorch```
@@ -175,3 +183,6 @@
 
+### [EfficientNet](https://arxiv.org/abs/1905.11946)
+Defaults: (layer = 1 from features, layer_output_size = 1280 for efficientnet_b0 model)<br>
 
+
 ## To-do
@@ -178,0 +189,0 @@ - Benchmark speed and accuracy

img2vec_pytorch/img_to_vec.py

@@ -13,5 +13,16 @@ import torch
         'resnet101': 2048,
-        'resnet152': 2048,
+        'resnet152': 2048
     }
 
+    EFFICIENTNET_OUTPUT_SIZES = {
+        'efficientnet_b0': 1280,
+        'efficientnet_b1': 1280,
+        'efficientnet_b2': 1408,
+        'efficientnet_b3': 1536,
+        'efficientnet_b4': 1792,
+        'efficientnet_b5': 2048,
+        'efficientnet_b6': 2304,
+        'efficientnet_b7': 2560
+    }
+
     def __init__(self, cuda=False, model='resnet-18', layer='default', layer_output_size=512):
@@ -50,3 +61,3 @@ """ Img2Vec
                 my_embedding = torch.zeros(len(img), self.layer_output_size)
-            elif self.model_name == 'densenet':
+            elif self.model_name == 'densenet' or 'efficientnet' in self.model_name:
                 my_embedding = torch.zeros(len(img), self.layer_output_size, 7, 7)
@@ -69,3 +80,3 @@ else:
                     return my_embedding.numpy()[:, :]
-                elif self.model_name == 'densenet':
+                elif self.model_name == 'densenet' or 'efficientnet' in self.model_name:
                     return torch.mean(my_embedding, (2, 3), True).numpy()[:, :, 0, 0]
@@ -79,3 +90,3 @@ else:
                 my_embedding = torch.zeros(1, self.layer_output_size)
-            elif self.model_name == 'densenet':
+            elif self.model_name == 'densenet' or 'efficientnet' in self.model_name:
                 my_embedding = torch.zeros(1, self.layer_output_size, 7, 7)
@@ -160,3 +171,32 @@ else:
 
+        elif "efficientnet" in model_name:
+            # efficientnet-b0 ~ efficientnet-b7
+            if model_name == "efficientnet_b0":
+                model = models.efficientnet_b0(pretrained=True)
+            elif model_name == "efficientnet_b1":
+                model = models.efficientnet_b1(pretrained=True)
+            elif model_name == "efficientnet_b2":
+                model = models.efficientnet_b2(pretrained=True)
+            elif model_name == "efficientnet_b3":
+                model = models.efficientnet_b3(pretrained=True)
+            elif model_name == "efficientnet_b4":
+                model = models.efficientnet_b4(pretrained=True)
+            elif model_name == "efficientnet_b5":
+                model = models.efficientnet_b5(pretrained=True)
+            elif model_name == "efficientnet_b6":
+                model = models.efficientnet_b6(pretrained=True)
+            elif model_name == "efficientnet_b7":
+                model = models.efficientnet_b7(pretrained=True)
+            else:
+                raise KeyError('Un support %s.' % model_name)
+
+            if layer == 'default':
+                layer = model.features
+                self.layer_output_size = self.EFFICIENTNET_OUTPUT_SIZES[model_name]
+            else:
+                raise KeyError('Un support %s for layer parameters' % model_name)
+
+            return model, layer
+
         else:
             raise KeyError('Model %s was not found' % model_name)

img2vec_pytorch/test_img_to_vec.py

@@ -5,3 +5,2 @@ import unittest
 
-
 from . import Img2Vec
@@ -12,3 +11,3 @@
         img2vec = Img2Vec()
-        img = Image.open('./example/test_images/cat.jpg')
+        img = Image.open('./example/test_images/cat.jpg').convert('RGB')
         vec = img2vec.get_vec(img)
@@ -21,3 +20,3 @@ self.assertEqual(True, isinstance(vec, numpy.ndarray))
         img2vec = Img2Vec(model='alexnet')
-        img = Image.open('./example/test_images/cat.jpg')
+        img = Image.open('./example/test_images/cat.jpg').convert('RGB')
         vec = img2vec.get_vec(img)
@@ -30,3 +29,3 @@ self.assertEqual(True, isinstance(vec, numpy.ndarray))
         img2vec = Img2Vec(model='vgg')
-        img = Image.open('./example/test_images/cat.jpg')
+        img = Image.open('./example/test_images/cat.jpg').convert('RGB')
         vec = img2vec.get_vec(img)
@@ -39,3 +38,3 @@ self.assertEqual(True, isinstance(vec, numpy.ndarray))
         img2vec = Img2Vec(model='densenet')
-        img = Image.open('./example/test_images/cat.jpg')
+        img = Image.open('./example/test_images/cat.jpg').convert('RGB')
         vec = img2vec.get_vec(img)
@@ -46,3 +45,67 @@ self.assertEqual(True, isinstance(vec, numpy.ndarray))
 
+    def test_efficientnet_b0(self):
+        img2vec = Img2Vec(model='efficientnet_b0')
+        img = Image.open('./example/test_images/cat.jpg')
+        vec = img2vec.get_vec(img)
+        self.assertEqual(True, isinstance(vec, numpy.ndarray))
+        self.assertEqual(1, vec.ndim)
+        self.assertEqual(1280, vec.size)
+
+    def test_efficientnet_b1(self):
+        img2vec = Img2Vec(model='efficientnet_b1')
+        img = Image.open('./example/test_images/cat.jpg')
+        vec = img2vec.get_vec(img)
+        self.assertEqual(True, isinstance(vec, numpy.ndarray))
+        self.assertEqual(1, vec.ndim)
+        self.assertEqual(1280, vec.size)
+
+    def test_efficientnet_b2(self):
+        img2vec = Img2Vec(model='efficientnet_b2')
+        img = Image.open('./example/test_images/cat.jpg')
+        vec = img2vec.get_vec(img)
+        self.assertEqual(True, isinstance(vec, numpy.ndarray))
+        self.assertEqual(1, vec.ndim)
+        self.assertEqual(1408, vec.size)
+
+    def test_efficientnet_b3(self):
+        img2vec = Img2Vec(model='efficientnet_b3')
+        img = Image.open('./example/test_images/cat.jpg')
+        vec = img2vec.get_vec(img)
+        self.assertEqual(True, isinstance(vec, numpy.ndarray))
+        self.assertEqual(1, vec.ndim)
+        self.assertEqual(1536, vec.size)
+
+    def test_efficientnet_b4(self):
+        img2vec = Img2Vec(model='efficientnet_b4')
+        img = Image.open('./example/test_images/cat.jpg')
+        vec = img2vec.get_vec(img)
+        self.assertEqual(True, isinstance(vec, numpy.ndarray))
+        self.assertEqual(1, vec.ndim)
+        self.assertEqual(1792, vec.size)
+
+    def test_efficientnet_b5(self):
+        img2vec = Img2Vec(model='efficientnet_b5')
+        img = Image.open('./example/test_images/cat.jpg')
+        vec = img2vec.get_vec(img)
+        self.assertEqual(True, isinstance(vec, numpy.ndarray))
+        self.assertEqual(1, vec.ndim)
+        self.assertEqual(2048, vec.size)
+
+    def test_efficientnet_b6(self):
+        img2vec = Img2Vec(model='efficientnet_b6')
+        img = Image.open('./example/test_images/cat.jpg')
+        vec = img2vec.get_vec(img)
+        self.assertEqual(True, isinstance(vec, numpy.ndarray))
+        self.assertEqual(1, vec.ndim)
+        self.assertEqual(2304, vec.size)
+
+    def test_efficientnet_b7(self):
+        img2vec = Img2Vec(model='efficientnet_b7')
+        img = Image.open('./example/test_images/cat.jpg')
+        vec = img2vec.get_vec(img)
+        self.assertEqual(True, isinstance(vec, numpy.ndarray))
+        self.assertEqual(1, vec.ndim)
+        self.assertEqual(2560, vec.size)
+
 if __name__ == "__main__":
     unittest.main()

PKG-INFO

 Metadata-Version: 2.1
 Name: img2vec_pytorch
-Version: 1.0.0
+Version: 1.0.1
 Summary: Use pre-trained models in PyTorch to extract vector embeddings for any image
@@ -30,17 +30,25 @@ Home-page: https://github.com/christiansafka/img2vec
 ## Available models
- - Resnet-18 (CPU, GPU)
- 	- Returns vector length 512
- - Alexnet (CPU, GPU)
- 	- Returns vector length 4096
- - Vgg-11 (CPU, GPU)
-	- Returns vector length 4096
- - Densenet (CPU, GPU)
-	- Returns vector length 1024
+|Model name|Return vector length|
+|----|----|
+|Resnet-18|512|
+|Alexnet|4096|
+|Vgg-11|4096|
+|Densenet|1024|
+|efficientnet_b0|1280|
+|efficientnet_b1|1280|
+|efficientnet_b2|1408|
+|efficientnet_b3|1536|
+|efficientnet_b4|1792|
+|efficientnet_b5|2048|
+|efficientnet_b6|2304|
+|efficientnet_b7|2560|
 
 ## Installation
 
-Tested on Python 3.6
+Tested on Python 3.6 and torchvision 0.11.0 (nightly, 2021-09-25) 
 
 Requires Pytorch: http://pytorch.org/
 
+```conda install -c pytorch-nightly torchvision```
+
 ```pip install img2vec_pytorch```
@@ -175,3 +183,6 @@
 
+### [EfficientNet](https://arxiv.org/abs/1905.11946)
+Defaults: (layer = 1 from features, layer_output_size = 1280 for efficientnet_b0 model)<br>
 
+
 ## To-do
@@ -178,0 +189,0 @@ - Benchmark speed and accuracy

README.md

@@ -14,17 +14,25 @@ # Image 2 Vec with PyTorch
 ## Available models
- - Resnet-18 (CPU, GPU)
- 	- Returns vector length 512
- - Alexnet (CPU, GPU)
- 	- Returns vector length 4096
- - Vgg-11 (CPU, GPU)
-	- Returns vector length 4096
- - Densenet (CPU, GPU)
-	- Returns vector length 1024
+|Model name|Return vector length|
+|----|----|
+|Resnet-18|512|
+|Alexnet|4096|
+|Vgg-11|4096|
+|Densenet|1024|
+|efficientnet_b0|1280|
+|efficientnet_b1|1280|
+|efficientnet_b2|1408|
+|efficientnet_b3|1536|
+|efficientnet_b4|1792|
+|efficientnet_b5|2048|
+|efficientnet_b6|2304|
+|efficientnet_b7|2560|
 
 ## Installation
 
-Tested on Python 3.6
+Tested on Python 3.6 and torchvision 0.11.0 (nightly, 2021-09-25) 
 
 Requires Pytorch: http://pytorch.org/
 
+```conda install -c pytorch-nightly torchvision```
+
 ```pip install img2vec_pytorch```
@@ -159,3 +167,6 @@
 
+### [EfficientNet](https://arxiv.org/abs/1905.11946)
+Defaults: (layer = 1 from features, layer_output_size = 1280 for efficientnet_b0 model)<br>
 
+
 ## To-do
@@ -162,0 +173,0 @@ - Benchmark speed and accuracy

setup.py

@@ -10,3 +10,3 @@ from setuptools import setup, find_packages
     name='img2vec_pytorch',
-    version='1.0.0',
+    version='1.0.1',
     description='Use pre-trained models in PyTorch to extract vector embeddings for any image',
@@ -13,0 +13,0 @@ long_description=long_description,

Alert delta unavailable

Currently unable to show alert delta for PyPI packages.

Improved metrics

Total package byte prevSize

31855

20.43%

Lines of code

285

47.67%