search_image_example.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "import urllib2\n",
    "import os\n",
    "import sklearn.neighbors\n",
    "import itertools\n",
    "import IPython\n",
    "import numpy\n",
    "\n",
    "import chainer.serializers\n",
    "\n",
    "import illust2comment\n",
    "import illust2comment.model\n",
    "import illust2comment.utility"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# arguments\n",
    "# 実際のパスで置き換えてください\n",
    "GPU = -1\n",
    "IMAGE_MODEL_PATH = \"data/nico_illust_tag_v2.hdf5\"\n",
    "COMMENT_DATA = \"data/seiga_comment_random.tsv\"\n",
    "\n",
    "IMAGES_DIRECTORY = \"image-data/\"\n",
    "QUERY_IMAGE_URL = \"http://3d.nicovideo.jp/alicia/img/profile_character.png\"\n",
    "\n",
    "TARGET_IMAGES_NUM = 1000"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "image_model = illust2comment.model.ImageModel(406)\n",
    "chainer.serializers.load_hdf5(IMAGE_MODEL_PATH, image_model.functions)\n",
    "if GPU >= 0:\n",
    "    chainer.cuda.check_cuda_available()\n",
    "    chainer.cuda.get_device(GPU).use()\n",
    "    xp = chainer.cuda.cupy\n",
    "    image_model.functions.to_gpu()\n",
    "else:\n",
    "    xp = numpy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "class ImageFeatureExtractor(object):\n",
    "    def __init__(self, image_model, xp):\n",
    "        self.xp = xp\n",
    "        self.image_model = image_model\n",
    "\n",
    "    def get_image_feature_from_url(self, image_url):\n",
    "        return self.get_image_feature(urllib2.urlopen(image_url))\n",
    "\n",
    "    def get_image_feature(self, image_path):\n",
    "        img_array = self.xp.array(illust2comment.utility.img2array(illust2comment.utility.load_image(image_path)))\n",
    "        return chainer.cuda.to_cpu(image_model.feature(img_array).data)[0]\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "features = []\n",
    "image_pathes = []\n",
    "image_feature_extractor = ImageFeatureExtractor(image_model, xp)\n",
    "n = 0\n",
    "import time\n",
    "print(time.time())\n",
    "for content_id, comment in illust2comment.utility.load_id_comments(COMMENT_DATA):\n",
    "    image_path = os.path.join(IMAGES_DIRECTORY, \"{}.jpg\".format(content_id[2:]))\n",
    "    if not os.path.exists(image_path):\n",
    "        continue\n",
    "    if image_path in image_pathes:\n",
    "        continue\n",
    "        \n",
    "    n += 1\n",
    "    feature = image_feature_extractor.get_image_feature(image_path)\n",
    "    features.append(feature)\n",
    "    image_pathes.append(image_path)\n",
    "    if n % 100 == 0:\n",
    "        print(n)\n",
    "    if n > TARGET_IMAGES_NUM:\n",
    "        break\n",
    "print(time.time())\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "engine = sklearn.neighbors.NearestNeighbors(n_neighbors=10, algorithm='auto', metric='euclidean')\n",
    "engine.fit(features)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "image_url = QUERY_IMAGE_URL\n",
    "IPython.display.display(IPython.display.Image(url=image_url))\n",
    "feature = image_feature_extractor.get_image_feature_from_url(image_url)\n",
    "distances_batch, target_indices_batch = engine.kneighbors([feature])\n",
    "for distance, target_index in itertools.izip(distances_batch[0], target_indices_batch[0]):\n",
    "    IPython.display.display(IPython.display.Image(filename=image_pathes[target_index]))\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 2",
   "language": "python",
   "name": "python2"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}