Use pre-trained models in PyTorch to extract vector embeddings for any image
Project description
Image 2 Vec with PyTorch
Medium post on building the first version from scratch: https://becominghuman.ai/extract-a-feature-vector-for-any-image-with-pytorch-9717561d1d4c
Looking for a simpler image vector integration for your project? Check out our free API at https://latentvector.space
Applications of image embeddings:
- Ranking for recommender systems
- Clustering images to different categories
- Classification tasks
- Image compression
Available models
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 and torchvision 0.11.0 (nightly, 2021-09-25)
Requires Pytorch: http://pytorch.org/
conda install -c pytorch-nightly torchvision
pip install img2vec_pytorch
Run test
python -m img2vec_pytorch.test_img_to_vec
Using img2vec as a library
from img2vec_pytorch import Img2Vec
from PIL import Image
# Initialize Img2Vec with GPU
img2vec = Img2Vec(cuda=True)
# Read in an image (rgb format)
img = Image.open('test.jpg')
# Get a vector from img2vec, returned as a torch FloatTensor
vec = img2vec.get_vec(img, tensor=True)
# Or submit a list
vectors = img2vec.get_vec(list_of_PIL_images)
For running the example, you will additionally need:
- Pillow:
pip install Pillow
- Sklearn
pip install scikit-learn
Running the example
git clone https://github.com/christiansafka/img2vec.git
cd img2vec/example
python test_img_similarity.py
Expected output
Which filename would you like similarities for?
cat.jpg
0.72832 cat2.jpg
0.641478 catdog.jpg
0.575845 face.jpg
0.516689 face2.jpg
Which filename would you like similarities for?
face2.jpg
0.668525 face.jpg
0.516689 cat.jpg
0.50084 cat2.jpg
0.484863 catdog.jpg
Try adding your own photos!
Img2Vec Params
cuda = (True, False) # Run on GPU? default: False
model = ('resnet-18', 'alexnet', 'vgg', 'densenet') # Which model to use? default: 'resnet-18'
Advanced users
Read only file systems
If you use this library from the app running in read only environment (for example, docker container), specify writable directory where app can store pre-trained models.
export TORCH_HOME=/tmp/torch
Additional Parameters
layer = 'layer_name' or int # For advanced users, which layer of the model to extract the output from. default: 'avgpool'
layer_output_size = int # Size of the output of your selected layer
Resnet-18
Defaults: (layer = 'avgpool', layer_output_size = 512)
Layer parameter must be an string representing the name of a layer below
conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3, bias=False)
bn1 = nn.BatchNorm2d(64)
relu = nn.ReLU(inplace=True)
maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
layer1 = self._make_layer(block, 64, layers[0])
layer2 = self._make_layer(block, 128, layers[1], stride=2)
layer3 = self._make_layer(block, 256, layers[2], stride=2)
layer4 = self._make_layer(block, 512, layers[3], stride=2)
avgpool = nn.AvgPool2d(7)
fc = nn.Linear(512 * block.expansion, num_classes)
Alexnet
Defaults: (layer = 2, layer_output_size = 4096)
Layer parameter must be an integer representing one of the layers below
alexnet.classifier = nn.Sequential(
7. nn.Dropout(), < - output_size = 9216
6. nn.Linear(256 * 6 * 6, 4096), < - output_size = 4096
5. nn.ReLU(inplace=True), < - output_size = 4096
4. nn.Dropout(), < - output_size = 4096
3. nn.Linear(4096, 4096), < - output_size = 4096
2. nn.ReLU(inplace=True), < - output_size = 4096
1. nn.Linear(4096, num_classes), < - output_size = 4096
)
Vgg
Defaults: (layer = 2, layer_output_size = 4096)
vgg.classifier = nn.Sequential(
nn.Linear(512 * 7 * 7, 4096),
nn.ReLU(True),
nn.Dropout(),
nn.Linear(4096, 4096),
nn.ReLU(True),
nn.Dropout(),
nn.Linear(4096, num_classes),
)
Densenet
Defaults: (layer = 1 from features, layer_output_size = 1024)
densenet.features = nn.Sequential(OrderedDict([
('conv0', nn.Conv2d(3, num_init_features, kernel_size=7, stride=2,
padding=3, bias=False)),
('norm0', nn.BatchNorm2d(num_init_features)),
('relu0', nn.ReLU(inplace=True)),
('pool0', nn.MaxPool2d(kernel_size=3, stride=2, padding=1)),
]))
EfficientNet
Defaults: (layer = 1 from features, layer_output_size = 1280 for efficientnet_b0 model)
To-do
- Benchmark speed and accuracy
- Add ability to fine-tune on input data
- Export documentation to a normal place
Project details
Release history Release notifications | RSS feed
Download files
Download the file for your platform. If you're not sure which to choose, learn more about installing packages.
Source Distribution
Built Distribution
File details
Details for the file img2vec_pytorch-1.0.1.tar.gz
.
File metadata
- Download URL: img2vec_pytorch-1.0.1.tar.gz
- Upload date:
- Size: 6.4 kB
- Tags: Source
- Uploaded using Trusted Publishing? No
- Uploaded via: twine/3.7.1 importlib_metadata/4.10.0 pkginfo/1.8.2 requests/2.26.0 requests-toolbelt/0.9.1 tqdm/4.62.3 CPython/3.9.9
File hashes
Algorithm | Hash digest | |
---|---|---|
SHA256 | c173b2ef8b4de673017f7cbb90d7febf48403c28d01aed47a231c3c5316d24f1 |
|
MD5 | d80670e9afea35c390824745b328fd31 |
|
BLAKE2b-256 | 0b3f68566abae6d5f49ec3180e7402d65be94678561fd3a1f406826178cbe808 |
File details
Details for the file img2vec_pytorch-1.0.1-py3-none-any.whl
.
File metadata
- Download URL: img2vec_pytorch-1.0.1-py3-none-any.whl
- Upload date:
- Size: 6.9 kB
- Tags: Python 3
- Uploaded using Trusted Publishing? No
- Uploaded via: twine/3.7.1 importlib_metadata/4.10.0 pkginfo/1.8.2 requests/2.26.0 requests-toolbelt/0.9.1 tqdm/4.62.3 CPython/3.9.9
File hashes
Algorithm | Hash digest | |
---|---|---|
SHA256 | b68cf9a30c8c17b0c19279f537e1dce6c78370e191414bf6589f62ae9c140adf |
|
MD5 | 1c99cbbf2efc90de1a44b0cb8fe8c46a |
|
BLAKE2b-256 | 8a1230eba1aa241abf034b7dfc4f76c7b95f9d0bc8fa374fc4d69fb4f095e0be |