add images to readme

Author: thomwolf

README.md

@@ -48,7 +48,7 @@ See the [doc section](#doc) below for details on these classes and methods.
 
 ```python
 import torch
-from pytorch_pretrained_biggan import (BigGAN, one_hot_from_name, truncated_noise_sample,
+from pytorch_pretrained_biggan import (BigGAN, one_hot_from_names, truncated_noise_sample,
                                        save_as_images, display_in_terminal)
 
 # OPTIONAL: if you want to have more information on what's happening, activate the logger as follows
@@ -58,26 +58,30 @@ logging.basicConfig(level=logging.INFO)
 # Load pre-trained model tokenizer (vocabulary)
 model = BigGAN.from_pretrained('biggan-deep-256')
 
-# Prepare a dogball input
+# Prepare a input
 truncation = 0.4
-class_vector = (one_hot_from_name('tennis ball') + one_hot_from_name('chihuahua')) / 2
-noise_vector = truncated_noise_sample(truncation=truncation)
+class_vector = one_hot_from_names(['soap bubble', 'coffee', 'mushroom'], batch_size=3)
+noise_vector = truncated_noise_sample(truncation=truncation, batch_size=3)
 
 # All in tensors
 noise_vector = torch.from_numpy(noise_vector)
 class_vector = torch.from_numpy(class_vector)
 
-# Generate a dog ball
-dogball = model(noise_vector, class_vector, truncation)
-
-# Save results as png images
-save_as_images(dogball)
+# Generate an image
+output = model(noise_vector, class_vector, truncation)
 
 # If you have a sixtel compatible terminal you can display the images in the terminal
 # (see https://github.com/saitoha/libsixel for details)
-display_in_terminal(dogball)
+display_in_terminal(output)
+
+# Save results as png images
+save_as_images(output)
 ```
 
+![output_0](assets/output_0.png)
+![output_1](assets/output_1.png)
+![output_2](assets/output_2.png)
+
 ## Doc
 
 ### Loading DeepMind's pre-trained weights
@@ -190,7 +194,7 @@ Here are some details on these methods:
     - Output:
         - array of shape (batch_size, 1000)
 
-- `one_hot_from_name(class_name, batch_size=1)`:
+- `one_hot_from_names(class_name, batch_size=1)`:
 
     Create a one-hot vector from the name of an imagenet class ('tennis ball', 'daisy', ...). We use NLTK's wordnet search to try to find the relevant synset of ImageNet and take the first one. If we can't find it direcly, we look at the hyponyms and hypernyms of the class name.
     - Params:

assets/output_0.png

@@ -1,5 +1,5 @@
 tensorflow
 tensorflow-hub
 Pillow
-NLTK
-python-libsixel
+nltk
+libsixel-python

assets/output_1.png

@@ -3,4 +3,4 @@
 from .file_utils import PYTORCH_PRETRAINED_BIGGAN_CACHE, cached_path
 from .utils import (truncated_noise_sample, save_as_images,
                     convert_to_images, display_in_terminal,
-                    one_hot_from_int, one_hot_from_name)
+                    one_hot_from_int, one_hot_from_names)

assets/output_2.png

@@ -298,9 +298,8 @@ def forward(self, z, class_label, truncation):
 
 if __name__ == "__main__":
     import PIL
-    from .utils import truncated_noise_sample, save_as_images
+    from .utils import truncated_noise_sample, save_as_images, one_hot_from_names
     from .convert_tf_to_pytorch import load_tf_weights_in_biggan
-    from .imagenet import one_hot_from_name
 
     load_cache = False
     cache_path = './saved_model.pt'
@@ -316,7 +315,7 @@ def forward(self, z, class_label, truncation):
 
     truncation = 0.4
     noise = truncated_noise_sample(batch_size=2, truncation=truncation)
-    label = one_hot_from_name('diver', batch_size=2)
+    label = one_hot_from_names('diver', batch_size=2)
 
     # Tests
     # noise = np.zeros((1, 128))