From 1a4e1702ec09c154c95fb03ac23de47ff3dc8bcd Mon Sep 17 00:00:00 2001
From: Varun Suresh <fab.varun@gmail.com>
Date: Tue, 28 Mar 2017 23:16:55 -0500
Subject: [PATCH] Added a demo for caffe and also updated the README

---
 README.md     | 33 ++++++++++++++++++++++++++++++++-
 demo_caffe.py | 43 +++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 75 insertions(+), 1 deletion(-)
 create mode 100644 demo_caffe.py

diff --git a/README.md b/README.md
index 4d435b90..d46a330f 100644
--- a/README.md
+++ b/README.md
@@ -15,7 +15,7 @@ For details, please see [our project page](http://illustration2vec.net) or
   several pre-trained models from http://illustration2vec.net,
   or execute ``get_models.sh`` in this repository).
 * ``numpy`` and ``scipy``
-* ``PIL`` (Python Imaging Library) or its alternatives (e.g., ``Pillow``) 
+* ``PIL`` (Python Imaging Library) or its alternatives (e.g., ``Pillow``)
 * ``skimage`` (Image processing library for python)
 
 In addition to the above libraries and the pre-trained models, `i2v` requires
@@ -149,6 +149,37 @@ shape: (1, 512), dtype: uint8
    42 181  38 254 177 232 150  99]]
 ```
 
+# Caffe Demo
+
+If you have caffe installed, use the `demo_caffe.py` to run a demo. For help,
+type python `demo_caffe.py --help`. You need to specify the directory where the
+model definition, model and the tags are stored.
+
+Assuming you have caffe setup correctly and all the files related to the model
+stored in the `models` directory, run `python demo_caffe.py` and you should see
+the following output :
+
+```
+[{'rating': [(u'safe', 0.978573203086853), (u'questionable', 0.02053508535027504),
+(u'explicit', 0.0006299669039435685)], 'character': [(u'hatsune miku', 0.999999463558197),
+(u'hatsune miku (append)', 0.000614884658716619), (u'kaito', 0.00040319858817383647),
+(u'hiiragi kagami', 0.00026519616949371994), (u'kagamine rin', 0.00023899678490124643),
+(u'megurine luka', 8.804832032183185e-05), (u'sakura miku', 4.707770131062716e-05),
+(u'kagamine len', 4.015495505882427e-05), (u'meiko', 2.7492344088386744e-05),
+(u'yuki miku', 1.8482929590390995e-05)], 'copyright': [(u'vocaloid', 0.9999998807907104),
+(u'project diva (series)', 0.0013757868437096477), (u'pangya', 0.00047354138223454356),
+(u'lucky star', 0.00034585013054311275), (u'project diva 2nd', 0.00028244982240721583),
+(u'pixiv', 9.411591599928215e-05), (u'utau', 9.144698560703546e-05),
+(u'transformers', 8.012886974029243e-05), (u'world is mine (vocaloid)', 7.95463056419976e-05),
+(u'macross', 5.722872447222471e-05)], 'general': [(u'thighhighs', 0.9956372380256653),
+(u'1girl', 0.9873462319374084), (u'twintails', 0.9812834858894348),
+(u'solo', 0.9632901549339294), (u'aqua hair', 0.9167951345443726),
+(u'long hair', 0.8817108273506165), (u'very long hair', 0.8326570987701416),
+(u'detached sleeves', 0.7448859810829163), (u'skirt', 0.6780790090560913),
+(u'necktie', 0.5608367919921875)]}]
+
+```
+
 # License
 The pre-trained models and the other files we have provided are licensed
 under the MIT License.
diff --git a/demo_caffe.py b/demo_caffe.py
new file mode 100644
index 00000000..c2751142
--- /dev/null
+++ b/demo_caffe.py
@@ -0,0 +1,43 @@
+# File to test out the illustration2vec file :
+
+import caffe
+import os
+from i2v.caffe_i2v import make_i2v_with_caffe
+from PIL import Image
+import argparse
+
+parser = argparse.ArgumentParser()
+
+parser.add_argument('--img_path',
+					type=str,
+					required=False,
+					help='Enter the path to the image',
+					default='images/miku.jpg'
+					)
+
+parser.add_argument('--model_dir',
+					type=str,
+					required=False,
+					help='Enter the path to the models folder',
+					default='models/'
+					)
+args = parser.parse_args()
+
+# If you have caffe with the gpu installed, use this. Otherwise comment
+# the line below.
+caffe.set_mode_gpu()
+
+# Download the model file, prototext and the tag list using the script
+# here - https://github.com/rezoo/illustration2vec/get_models.sh and
+# store it in a directory called models.
+
+model_def = os.path.join(args.model_dir,'illust2vec_tag.prototxt')
+model_wieghts = os.path.join(args.model_dir,'illust2vec_tag_ver200.caffemodel')
+tag_path = os.path.join(args.model_dir,'tag_list.json')
+
+net = make_i2v_with_caffe(model_def,model_wieghts,tag_path=tag_path)
+
+img = Image.open(args.img_path)
+
+print net.estimate_top_tags([img])
+# Take a look at i2v.base.py for other methods for the "caffe_i2v" class.