├── .dockerignore ├── .github └── workflows │ ├── docker.yml │ └── main.yml ├── .gitignore ├── Dockerfile ├── Dockerfile.gpu ├── LICENSE ├── README.md ├── docker-compose.gpu.yml ├── docker-compose.yml ├── examples ├── index.py ├── save_index.py ├── storage.py └── yolo_detector.py ├── images ├── 000000055167.jpg ├── 000000084362.jpg ├── 000000101068.jpg ├── 000000134112.jpg ├── 000000157138.jpg ├── 000000170545.jpg ├── 000000263463.jpg ├── 000000275727.jpg ├── 000000304812.jpg ├── 000000306582.jpg ├── 000000327617.jpg ├── 000000404805.jpg ├── 000000468124.jpg ├── 000000482735.jpg └── 000000534041.jpg ├── imsearch ├── __init__.py ├── backend │ ├── __init__.py │ ├── extractor.py │ ├── feature_extractor.py │ └── object_detector │ │ ├── __init__.py │ │ └── yolo │ │ ├── __init__.py │ │ ├── detector.py │ │ ├── models.py │ │ └── utils.py ├── config.py ├── exception.py ├── extractor.py ├── index.py ├── nmslib.py ├── repository │ ├── __init__.py │ └── mongo.py ├── storage │ ├── __init__.py │ ├── abstract.py │ ├── aws.py │ └── gcloud.py └── utils │ ├── __init__.py │ └── image.py ├── install.sh ├── requirements.txt ├── setup.py └── tests ├── test_index.py └── test_init.py /.dockerignore: -------------------------------------------------------------------------------- 1 | .git 2 | .git/** 3 | .github/ 4 | .github/** 5 | .gitignore 6 | .dockerignore 7 | *.yml 8 | # *.md 9 | Dockerfile 10 | Dockerfile.gpu 11 | -------------------------------------------------------------------------------- /.github/workflows/docker.yml: -------------------------------------------------------------------------------- 1 | name: Docker 2 | 3 | on: 4 | push: 5 | # Publish `master` as Docker `latest` image. 6 | branches: 7 | - master 8 | 9 | env: 10 | IMAGE_NAME: rmehta3/imsearch 11 | 12 | jobs: 13 | # Push image to GitHub Package Registry. 14 | # See also https://docs.docker.com/docker-hub/builds/ 15 | push: 16 | runs-on: ubuntu-latest 17 | if: github.event_name == 'push' 18 | 19 | steps: 20 | - uses: actions/checkout@v2 21 | 22 | - name: Build image 23 | run: docker build . --file Dockerfile --tag image 24 | 25 | - name: Log into registry 26 | run: echo ${{ secrets.DOCKER_PASSWORD}} | docker login -u ${{ secrets.DOCKER_USERNAME }} --password-stdin 27 | 28 | - name: Push image 29 | run: | 30 | IMAGE_ID=$IMAGE_NAME 31 | 32 | # Strip git ref prefix from version 33 | VERSION=$(echo "${{ github.ref }}" | sed -e 's,.*/\(.*\),\1,') 34 | 35 | # Strip "v" prefix from tag name 36 | [[ "${{ github.ref }}" == "refs/tags/"* ]] && VERSION=$(echo $VERSION | sed -e 's/^v//') 37 | 38 | # Use Docker `latest` tag convention 39 | [ "$VERSION" == "master" ] && VERSION=latest 40 | 41 | echo IMAGE_ID=$IMAGE_ID 42 | echo VERSION=$VERSION 43 | 44 | docker tag image $IMAGE_ID:$VERSION 45 | docker push $IMAGE_ID:$VERSION 46 | -------------------------------------------------------------------------------- /.github/workflows/main.yml: -------------------------------------------------------------------------------- 1 | # This is a basic workflow to help you get started with Actions 2 | 3 | name: CI 4 | 5 | # Controls when the action will run. Triggers the workflow on push or pull request 6 | # events but only for the master branch 7 | on: 8 | push: 9 | branches: [ master ] 10 | pull_request: 11 | branches: [ master ] 12 | 13 | # A workflow run is made up of one or more jobs that can run sequentially or in parallel 14 | jobs: 15 | # This workflow contains a single job called "build" 16 | build: 17 | 18 | runs-on: ubuntu-latest 19 | strategy: 20 | matrix: 21 | python-version: [3.6, 3.7, 3.8] 22 | 23 | # Steps represent a sequence of tasks that will be executed as part of the job 24 | steps: 25 | - uses: actions/checkout@v2 26 | - name: Set up Python ${{ matrix.python-version }} 27 | uses: actions/setup-python@v1 28 | with: 29 | python-version: ${{ matrix.python-version }} 30 | - name: Install dependencies 31 | run: | 32 | python -m pip install --upgrade pip 33 | pip install -r requirements.txt 34 | pip install --no-binary :all: nmslib 35 | - name: Test with nose2 36 | run: | 37 | pip install nose2 38 | nose2 39 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | # Byte-compiled / optimized / DLL files 2 | __pycache__/ 3 | *.py[cod] 4 | *$py.class 5 | 6 | # C extensions 7 | *.so 8 | 9 | # Distribution / packaging 10 | .Python 11 | build/ 12 | develop-eggs/ 13 | dist/ 14 | downloads/ 15 | eggs/ 16 | .eggs/ 17 | lib/ 18 | lib64/ 19 | parts/ 20 | sdist/ 21 | var/ 22 | wheels/ 23 | *.egg-info/ 24 | .installed.cfg 25 | *.egg 26 | MANIFEST 27 | 28 | # PyInstaller 29 | # Usually these files are written by a python script from a template 30 | # before PyInstaller builds the exe, so as to inject date/other infos into it. 31 | *.manifest 32 | *.spec 33 | 34 | # Installer logs 35 | pip-log.txt 36 | pip-delete-this-directory.txt 37 | 38 | # Unit test / coverage reports 39 | htmlcov/ 40 | .tox/ 41 | .coverage 42 | .coverage.* 43 | .cache 44 | nosetests.xml 45 | coverage.xml 46 | *.cover 47 | .hypothesis/ 48 | .pytest_cache/ 49 | 50 | # Translations 51 | *.mo 52 | *.pot 53 | 54 | # Django stuff: 55 | *.log 56 | local_settings.py 57 | db.sqlite3 58 | 59 | # Flask stuff: 60 | instance/ 61 | .webassets-cache 62 | 63 | # Scrapy stuff: 64 | .scrapy 65 | 66 | # Sphinx documentation 67 | docs/_build/ 68 | 69 | # PyBuilder 70 | target/ 71 | 72 | # Jupyter Notebook 73 | .ipynb_checkpoints 74 | 75 | # pyenv 76 | .python-version 77 | 78 | # celery beat schedule file 79 | celerybeat-schedule 80 | 81 | # SageMath parsed files 82 | *.sage.py 83 | 84 | # Environments 85 | .env 86 | .venv 87 | env/ 88 | venv/ 89 | ENV/ 90 | env.bak/ 91 | venv.bak/ 92 | .vscode/ 93 | 94 | # Spyder project settings 95 | .spyderproject 96 | .spyproject 97 | 98 | # Rope project settings 99 | .ropeproject 100 | 101 | # mkdocs documentation 102 | /site 103 | 104 | # mypy 105 | .mypy_cache/ 106 | 107 | # data, model weights & csv files 108 | data/ 109 | .pth 110 | .csv 111 | .log 112 | logs/ 113 | .config 114 | -------------------------------------------------------------------------------- /Dockerfile: -------------------------------------------------------------------------------- 1 | # ================================================================== 2 | # module list 3 | # ------------------------------------------------------------------ 4 | # python 3.6 (apt) 5 | # pytorch latest (pip) 6 | # ================================================================== 7 | 8 | FROM ubuntu:18.04 9 | ENV LANG C.UTF-8 10 | RUN APT_INSTALL="apt-get install -y --no-install-recommends" && \ 11 | PIP_INSTALL="python -m pip install --upgrade --no-cache-dir --retries 10 --timeout 60" && \ 12 | GIT_CLONE="git clone --depth 10" && \ 13 | \ 14 | rm -rf /var/lib/apt/lists/* \ 15 | /etc/apt/sources.list.d/cuda.list \ 16 | /etc/apt/sources.list.d/nvidia-ml.list && \ 17 | \ 18 | apt-get update && \ 19 | DEBIAN_FRONTEND=noninteractive $APT_INSTALL \ 20 | software-properties-common \ 21 | && \ 22 | add-apt-repository "deb http://security.ubuntu.com/ubuntu xenial-security main" && \ 23 | apt-get update && \ 24 | # ================================================================== 25 | # tools 26 | # ------------------------------------------------------------------ 27 | DEBIAN_FRONTEND=noninteractive $APT_INSTALL \ 28 | build-essential \ 29 | apt-utils \ 30 | ca-certificates \ 31 | wget \ 32 | libsm6 \ 33 | libxext6 \ 34 | libxrender-dev \ 35 | git \ 36 | bash \ 37 | nano \ 38 | vim \ 39 | libssl-dev \ 40 | curl \ 41 | unzip \ 42 | unrar \ 43 | && \ 44 | \ 45 | $GIT_CLONE https://github.com/Kitware/CMake ~/cmake && \ 46 | cd ~/cmake && \ 47 | ./bootstrap && \ 48 | make -j"$(nproc)" install && \ 49 | \ 50 | # ================================================================== 51 | # python 52 | # ------------------------------------------------------------------ 53 | \ 54 | DEBIAN_FRONTEND=noninteractive $APT_INSTALL \ 55 | software-properties-common \ 56 | && \ 57 | add-apt-repository ppa:deadsnakes/ppa && \ 58 | apt-get update && \ 59 | DEBIAN_FRONTEND=noninteractive $APT_INSTALL \ 60 | python3.6 \ 61 | python3.6-dev \ 62 | python3-distutils-extra \ 63 | && \ 64 | wget -O ~/get-pip.py \ 65 | https://bootstrap.pypa.io/get-pip.py && \ 66 | python3.6 ~/get-pip.py && \ 67 | ln -s /usr/bin/python3.6 /usr/local/bin/python3 && \ 68 | ln -s /usr/bin/python3.6 /usr/local/bin/python && \ 69 | $PIP_INSTALL \ 70 | setuptools \ 71 | && \ 72 | $PIP_INSTALL \ 73 | numpy \ 74 | scipy \ 75 | pandas \ 76 | cloudpickle \ 77 | scikit-image>=0.14.2 \ 78 | scikit-learn \ 79 | matplotlib \ 80 | Cython \ 81 | tqdm \ 82 | && \ 83 | \ 84 | # ================================================================== 85 | # pytorch 86 | # ------------------------------------------------------------------ 87 | \ 88 | $PIP_INSTALL \ 89 | future \ 90 | numpy \ 91 | protobuf \ 92 | enum34 \ 93 | pyyaml \ 94 | typing \ 95 | && \ 96 | $PIP_INSTALL \ 97 | --pre torch torchvision -f \ 98 | https://download.pytorch.org/whl/nightly/cpu/torch_nightly.html \ 99 | && \ 100 | \ 101 | # ================================================================== 102 | # config & cleanup 103 | # ------------------------------------------------------------------ 104 | \ 105 | ldconfig && \ 106 | apt-get clean && \ 107 | apt-get autoremove && \ 108 | rm -rf /var/lib/apt/lists/* /tmp/* ~/* 109 | 110 | # ================================================================== 111 | # install nmslib 112 | # ------------------------------------------------------------------ 113 | RUN pip install --no-binary :all: nmslib 114 | 115 | # ================================================================== 116 | # install imsearch 117 | # ------------------------------------------------------------------ 118 | WORKDIR /imsearch 119 | COPY . /imsearch 120 | 121 | # ================================================================== 122 | # test imsearch 123 | # ------------------------------------------------------------------ 124 | RUN python3 -m pip install --upgrade pip && \ 125 | pip install -r requirements.txt && \ 126 | pip install nose2 && \ 127 | nose2 128 | 129 | RUN pip install . 130 | 131 | CMD ["/bin/bash"] 132 | -------------------------------------------------------------------------------- /Dockerfile.gpu: -------------------------------------------------------------------------------- 1 | # ================================================================== 2 | # module list 3 | # ------------------------------------------------------------------ 4 | # python 3.6 (apt) 5 | # pytorch latest (pip) 6 | # ================================================================== 7 | 8 | FROM nvidia/cuda:10.1-cudnn7-devel-ubuntu18.04 9 | ENV LANG C.UTF-8 10 | RUN APT_INSTALL="apt-get install -y --no-install-recommends" && \ 11 | PIP_INSTALL="python -m pip install --upgrade --no-cache-dir --retries 10 --timeout 60" && \ 12 | GIT_CLONE="git clone --depth 10" && \ 13 | \ 14 | rm -rf /var/lib/apt/lists/* \ 15 | /etc/apt/sources.list.d/cuda.list \ 16 | /etc/apt/sources.list.d/nvidia-ml.list && \ 17 | \ 18 | apt-get update && \ 19 | DEBIAN_FRONTEND=noninteractive $APT_INSTALL \ 20 | software-properties-common \ 21 | && \ 22 | add-apt-repository "deb http://security.ubuntu.com/ubuntu xenial-security main" && \ 23 | apt-get update && \ 24 | \ 25 | # ================================================================== 26 | # tools 27 | # ------------------------------------------------------------------ 28 | \ 29 | DEBIAN_FRONTEND=noninteractive $APT_INSTALL \ 30 | build-essential \ 31 | apt-utils \ 32 | ca-certificates \ 33 | wget \ 34 | git \ 35 | vim \ 36 | libsm6 \ 37 | libxext6 \ 38 | libxrender-dev \ 39 | libssl-dev \ 40 | curl \ 41 | nano \ 42 | bash \ 43 | unzip \ 44 | unrar \ 45 | && \ 46 | \ 47 | $GIT_CLONE https://github.com/Kitware/CMake ~/cmake && \ 48 | cd ~/cmake && \ 49 | ./bootstrap && \ 50 | make -j"$(nproc)" install && \ 51 | \ 52 | # ================================================================== 53 | # python 54 | # ------------------------------------------------------------------ 55 | \ 56 | DEBIAN_FRONTEND=noninteractive $APT_INSTALL \ 57 | software-properties-common \ 58 | && \ 59 | add-apt-repository ppa:deadsnakes/ppa && \ 60 | apt-get update && \ 61 | DEBIAN_FRONTEND=noninteractive $APT_INSTALL \ 62 | python3.6 \ 63 | python3.6-dev \ 64 | python3-distutils-extra \ 65 | && \ 66 | wget -O ~/get-pip.py \ 67 | https://bootstrap.pypa.io/get-pip.py && \ 68 | python3.6 ~/get-pip.py && \ 69 | ln -s /usr/bin/python3.6 /usr/local/bin/python3 && \ 70 | ln -s /usr/bin/python3.6 /usr/local/bin/python && \ 71 | $PIP_INSTALL \ 72 | setuptools \ 73 | && \ 74 | $PIP_INSTALL \ 75 | numpy \ 76 | scipy \ 77 | pandas \ 78 | cloudpickle \ 79 | scikit-image>=0.14.2 \ 80 | scikit-learn \ 81 | matplotlib \ 82 | Cython \ 83 | tqdm \ 84 | && \ 85 | \ 86 | # ================================================================== 87 | # pytorch 88 | # ------------------------------------------------------------------ 89 | \ 90 | $PIP_INSTALL \ 91 | future \ 92 | numpy \ 93 | protobuf \ 94 | enum34 \ 95 | pyyaml \ 96 | typing \ 97 | && \ 98 | $PIP_INSTALL \ 99 | --pre torch torchvision -f \ 100 | https://download.pytorch.org/whl/nightly/cu101/torch_nightly.html \ 101 | && \ 102 | \ 103 | # ================================================================== 104 | # config & cleanup 105 | # ------------------------------------------------------------------ 106 | \ 107 | ldconfig && \ 108 | apt-get clean && \ 109 | apt-get autoremove && \ 110 | rm -rf /var/lib/apt/lists/* /tmp/* ~/* 111 | 112 | # ================================================================== 113 | # install nmslib 114 | # ------------------------------------------------------------------ 115 | RUN pip install --no-binary :all: nmslib 116 | 117 | # ================================================================== 118 | # install imsearch 119 | # ------------------------------------------------------------------ 120 | WORKDIR /imsearch 121 | COPY . /imsearch 122 | 123 | # ================================================================== 124 | # test imsearch 125 | # ------------------------------------------------------------------ 126 | RUN python3 -m pip install --upgrade pip && \ 127 | pip install -r requirements.txt && \ 128 | pip install nose2 && \ 129 | nose2 130 | 131 | RUN pip install . 132 | 133 | CMD ["/bin/bash"] 134 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | Apache License 2 | Version 2.0, January 2004 3 | http://www.apache.org/licenses/ 4 | 5 | TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION 6 | 7 | 1. Definitions. 8 | 9 | "License" shall mean the terms and conditions for use, reproduction, 10 | and distribution as defined by Sections 1 through 9 of this document. 11 | 12 | "Licensor" shall mean the copyright owner or entity authorized by 13 | the copyright owner that is granting the License. 14 | 15 | "Legal Entity" shall mean the union of the acting entity and all 16 | other entities that control, are controlled by, or are under common 17 | control with that entity. For the purposes of this definition, 18 | "control" means (i) the power, direct or indirect, to cause the 19 | direction or management of such entity, whether by contract or 20 | otherwise, or (ii) ownership of fifty percent (50%) or more of the 21 | outstanding shares, or (iii) beneficial ownership of such entity. 22 | 23 | "You" (or "Your") shall mean an individual or Legal Entity 24 | exercising permissions granted by this License. 25 | 26 | "Source" form shall mean the preferred form for making modifications, 27 | including but not limited to software source code, documentation 28 | source, and configuration files. 29 | 30 | "Object" form shall mean any form resulting from mechanical 31 | transformation or translation of a Source form, including but 32 | not limited to compiled object code, generated documentation, 33 | and conversions to other media types. 34 | 35 | "Work" shall mean the work of authorship, whether in Source or 36 | Object form, made available under the License, as indicated by a 37 | copyright notice that is included in or attached to the work 38 | (an example is provided in the Appendix below). 39 | 40 | "Derivative Works" shall mean any work, whether in Source or Object 41 | form, that is based on (or derived from) the Work and for which the 42 | editorial revisions, annotations, elaborations, or other modifications 43 | represent, as a whole, an original work of authorship. For the purposes 44 | of this License, Derivative Works shall not include works that remain 45 | separable from, or merely link (or bind by name) to the interfaces of, 46 | the Work and Derivative Works thereof. 47 | 48 | "Contribution" shall mean any work of authorship, including 49 | the original version of the Work and any modifications or additions 50 | to that Work or Derivative Works thereof, that is intentionally 51 | submitted to Licensor for inclusion in the Work by the copyright owner 52 | or by an individual or Legal Entity authorized to submit on behalf of 53 | the copyright owner. For the purposes of this definition, "submitted" 54 | means any form of electronic, verbal, or written communication sent 55 | to the Licensor or its representatives, including but not limited to 56 | communication on electronic mailing lists, source code control systems, 57 | and issue tracking systems that are managed by, or on behalf of, the 58 | Licensor for the purpose of discussing and improving the Work, but 59 | excluding communication that is conspicuously marked or otherwise 60 | designated in writing by the copyright owner as "Not a Contribution." 61 | 62 | "Contributor" shall mean Licensor and any individual or Legal Entity 63 | on behalf of whom a Contribution has been received by Licensor and 64 | subsequently incorporated within the Work. 65 | 66 | 2. Grant of Copyright License. Subject to the terms and conditions of 67 | this License, each Contributor hereby grants to You a perpetual, 68 | worldwide, non-exclusive, no-charge, royalty-free, irrevocable 69 | copyright license to reproduce, prepare Derivative Works of, 70 | publicly display, publicly perform, sublicense, and distribute the 71 | Work and such Derivative Works in Source or Object form. 72 | 73 | 3. Grant of Patent License. Subject to the terms and conditions of 74 | this License, each Contributor hereby grants to You a perpetual, 75 | worldwide, non-exclusive, no-charge, royalty-free, irrevocable 76 | (except as stated in this section) patent license to make, have made, 77 | use, offer to sell, sell, import, and otherwise transfer the Work, 78 | where such license applies only to those patent claims licensable 79 | by such Contributor that are necessarily infringed by their 80 | Contribution(s) alone or by combination of their Contribution(s) 81 | with the Work to which such Contribution(s) was submitted. If You 82 | institute patent litigation against any entity (including a 83 | cross-claim or counterclaim in a lawsuit) alleging that the Work 84 | or a Contribution incorporated within the Work constitutes direct 85 | or contributory patent infringement, then any patent licenses 86 | granted to You under this License for that Work shall terminate 87 | as of the date such litigation is filed. 88 | 89 | 4. Redistribution. You may reproduce and distribute copies of the 90 | Work or Derivative Works thereof in any medium, with or without 91 | modifications, and in Source or Object form, provided that You 92 | meet the following conditions: 93 | 94 | (a) You must give any other recipients of the Work or 95 | Derivative Works a copy of this License; and 96 | 97 | (b) You must cause any modified files to carry prominent notices 98 | stating that You changed the files; and 99 | 100 | (c) You must retain, in the Source form of any Derivative Works 101 | that You distribute, all copyright, patent, trademark, and 102 | attribution notices from the Source form of the Work, 103 | excluding those notices that do not pertain to any part of 104 | the Derivative Works; and 105 | 106 | (d) If the Work includes a "NOTICE" text file as part of its 107 | distribution, then any Derivative Works that You distribute must 108 | include a readable copy of the attribution notices contained 109 | within such NOTICE file, excluding those notices that do not 110 | pertain to any part of the Derivative Works, in at least one 111 | of the following places: within a NOTICE text file distributed 112 | as part of the Derivative Works; within the Source form or 113 | documentation, if provided along with the Derivative Works; or, 114 | within a display generated by the Derivative Works, if and 115 | wherever such third-party notices normally appear. The contents 116 | of the NOTICE file are for informational purposes only and 117 | do not modify the License. You may add Your own attribution 118 | notices within Derivative Works that You distribute, alongside 119 | or as an addendum to the NOTICE text from the Work, provided 120 | that such additional attribution notices cannot be construed 121 | as modifying the License. 122 | 123 | You may add Your own copyright statement to Your modifications and 124 | may provide additional or different license terms and conditions 125 | for use, reproduction, or distribution of Your modifications, or 126 | for any such Derivative Works as a whole, provided Your use, 127 | reproduction, and distribution of the Work otherwise complies with 128 | the conditions stated in this License. 129 | 130 | 5. Submission of Contributions. Unless You explicitly state otherwise, 131 | any Contribution intentionally submitted for inclusion in the Work 132 | by You to the Licensor shall be under the terms and conditions of 133 | this License, without any additional terms or conditions. 134 | Notwithstanding the above, nothing herein shall supersede or modify 135 | the terms of any separate license agreement you may have executed 136 | with Licensor regarding such Contributions. 137 | 138 | 6. Trademarks. This License does not grant permission to use the trade 139 | names, trademarks, service marks, or product names of the Licensor, 140 | except as required for reasonable and customary use in describing the 141 | origin of the Work and reproducing the content of the NOTICE file. 142 | 143 | 7. Disclaimer of Warranty. Unless required by applicable law or 144 | agreed to in writing, Licensor provides the Work (and each 145 | Contributor provides its Contributions) on an "AS IS" BASIS, 146 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or 147 | implied, including, without limitation, any warranties or conditions 148 | of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A 149 | PARTICULAR PURPOSE. You are solely responsible for determining the 150 | appropriateness of using or redistributing the Work and assume any 151 | risks associated with Your exercise of permissions under this License. 152 | 153 | 8. Limitation of Liability. In no event and under no legal theory, 154 | whether in tort (including negligence), contract, or otherwise, 155 | unless required by applicable law (such as deliberate and grossly 156 | negligent acts) or agreed to in writing, shall any Contributor be 157 | liable to You for damages, including any direct, indirect, special, 158 | incidental, or consequential damages of any character arising as a 159 | result of this License or out of the use or inability to use the 160 | Work (including but not limited to damages for loss of goodwill, 161 | work stoppage, computer failure or malfunction, or any and all 162 | other commercial damages or losses), even if such Contributor 163 | has been advised of the possibility of such damages. 164 | 165 | 9. Accepting Warranty or Additional Liability. While redistributing 166 | the Work or Derivative Works thereof, You may choose to offer, 167 | and charge a fee for, acceptance of support, warranty, indemnity, 168 | or other liability obligations and/or rights consistent with this 169 | License. However, in accepting such obligations, You may act only 170 | on Your own behalf and on Your sole responsibility, not on behalf 171 | of any other Contributor, and only if You agree to indemnify, 172 | defend, and hold each Contributor harmless for any liability 173 | incurred by, or claims asserted against, such Contributor by reason 174 | of your accepting any such warranty or additional liability. 175 | 176 | END OF TERMS AND CONDITIONS 177 | 178 | APPENDIX: How to apply the Apache License to your work. 179 | 180 | To apply the Apache License to your work, attach the following 181 | boilerplate notice, with the fields enclosed by brackets "[]" 182 | replaced with your own identifying information. (Don't include 183 | the brackets!) The text should be enclosed in the appropriate 184 | comment syntax for the file format. We also recommend that a 185 | file or class name and description of purpose be included on the 186 | same "printed page" as the copyright notice for easier 187 | identification within third-party archives. 188 | 189 | Copyright 2018 Riken Mehta 190 | 191 | Licensed under the Apache License, Version 2.0 (the "License"); 192 | you may not use this file except in compliance with the License. 193 | You may obtain a copy of the License at 194 | 195 | http://www.apache.org/licenses/LICENSE-2.0 196 | 197 | Unless required by applicable law or agreed to in writing, software 198 | distributed under the License is distributed on an "AS IS" BASIS, 199 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 200 | See the License for the specific language governing permissions and 201 | limitations under the License. 202 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # imSearch: A generic framework to build your own reverse image search engine 2 | 3 | ![CI](https://github.com/rikenmehta03/imsearch/workflows/CI/badge.svg?branch=master) 4 | 5 | imsearch helps to create your own custom, robust & scalable reverse image search engine. This project uses state of the art object detection algorithm ([yolov3](https://pjreddie.com/darknet/yolo/)) at its core to extract the features from an image. It uses an efficient cross-platform similarity search library [NMSLIB](https://github.com/nmslib/nmslib) for similarity search. [Redis](https://redis.io/) is used as a messaging queue between feature extractor and core engine. [MongoDB](https://www.mongodb.com/) is used to store the meta-data of all the indexed images. HD5 file system is used to store the feature vectors extracted from indexed images. 6 | 7 | ## Installation 8 | For the setup, a simple `install.sh` script can be used or can be installed using `pip`. 9 | Follow these simple steps to install imsearch library. 10 | - Feature extraction is GPU intensive process. So, to make the search real-time, running this engine on GPU enabled machine is recommended. 11 | - Install CUDA & NVIDIA graphics drivers ([here](https://medium.com/@taylordenouden/installing-tensorflow-gpu-on-ubuntu-18-04-89a142325138)) 12 | - Install `PyTorch` ([here](https://pytorch.org/get-started/locally/)) 13 | - Install `MongoDB` ([here](https://docs.mongodb.com/manual/tutorial/install-mongodb-on-ubuntu/)) 14 | - Install `Redis` ([here](https://www.digitalocean.com/community/tutorials/how-to-install-and-secure-redis-on-ubuntu-18-04)) 15 | - Run following commands 16 | ``` 17 | pip install --no-binary :all: nmslib 18 | pip install imsearch 19 | ``` 20 | 21 | ### Build from source using `install.sh` 22 | ``` 23 | git clone https://github.com/rikenmehta03/imsearch.git 24 | chmod +x install.sh 25 | ./install.sh 26 | ``` 27 | 28 | ## Example usage 29 | ``` 30 | import glob 31 | import imsearch 32 | 33 | all_images = glob.glob('path/to/image/folder') 34 | 35 | # Initialize the index 36 | index = imsearch.init('test') 37 | 38 | # Add single image to the index 39 | index.addImage(all_images[0]) 40 | 41 | # Add multiple image to the index 42 | index.addImageBatch(all_images[1:]) 43 | 44 | # Create index and make it ready for the search query 45 | index.createIndex() 46 | 47 | # find k nearest similar images 48 | # choose policy from 'object' or 'global'. Search results will change accordingly. 49 | # object: Object level matching. The engine will look for similarity at object level for every object detected in the image. 50 | # global: Overall similarity using single feature space on the whole image. 51 | similar = index.knnQuery('path/to/query/image', k=10, policy='object') 52 | ``` 53 | For detailed usage see [`examples/index.py`](examples/index.py) 54 | 55 | ## Docker 56 | If you don't have Docker/Docker-Compose check **Setup Docker** section 57 | 58 |
59 | Setup Docker 60 |

61 | 62 | ### Docker 63 | macOS: https://docs.docker.com/docker-for-mac/install/ 64 | 65 | linux: https://docs.docker.com/install/linux/docker-ce/ubuntu/ 66 | 67 | ### Docker Compose 68 | 69 | linux: https://docs.docker.com/compose/install/ 70 |

71 |
72 | 73 | ### For CPU 74 | ```bash 75 | docker-compose build 76 | docker-compose run imsearch 77 | ``` 78 | 79 | ### For GPU 80 | ```bash 81 | docker-compose -f docker-compose.gpu.yml build 82 | docker-compose -f docker-compose.gpu.yml run imsearch 83 | ``` 84 | 85 | ## Credit 86 | 87 | ### YOLOv3: An Incremental Improvement 88 | _Joseph Redmon, Ali Farhadi_
89 | 90 | **Abstract**
91 | We present some updates to YOLO! We made a bunch 92 | of little design changes to make it better. We also trained 93 | this new network that’s pretty swell. It’s a little bigger than 94 | last time but more accurate. It’s still fast though, don’t 95 | worry. At 320 × 320 YOLOv3 runs in 22 ms at 28.2 mAP, 96 | as accurate as SSD but three times faster. When we look 97 | at the old .5 IOU mAP detection metric YOLOv3 is quite 98 | good. It achieves 57.9 AP50 in 51 ms on a Titan X, compared 99 | to 57.5 AP50 in 198 ms by RetinaNet, similar performance 100 | but 3.8× faster. As always, all the code is online at 101 | https://pjreddie.com/yolo/. 102 | 103 | [[Paper]](https://pjreddie.com/media/files/papers/YOLOv3.pdf) [[Project Webpage]](https://pjreddie.com/darknet/yolo/) [[Authors' Implementation]](https://github.com/pjreddie/darknet) 104 | 105 | ``` 106 | @article{yolov3, 107 | title={YOLOv3: An Incremental Improvement}, 108 | author={Redmon, Joseph and Farhadi, Ali}, 109 | journal = {arXiv}, 110 | year={2018} 111 | } 112 | ``` 113 | 114 | ### PyTorch-YOLOv3 115 | Minimal PyTorch implementation of YOLOv3 [[GitHub]](https://github.com/eriklindernoren/PyTorch-YOLOv3) 116 | -------------------------------------------------------------------------------- /docker-compose.gpu.yml: -------------------------------------------------------------------------------- 1 | --- 2 | version: '3.7' 3 | services: 4 | 5 | imsearch: 6 | image: imsearch-gpu:latest 7 | container_name: imsearch-gpu 8 | build: 9 | context: . 10 | dockerfile: Dockerfile.gpu 11 | ports: 12 | - 5000:5000 13 | environment: 14 | - MONGO_URI=mongodb://mongodb:27017/ 15 | - REDIS_URI=redis://redis:6379/0 16 | - STORAGE_MODE=local 17 | depends_on: 18 | - redis 19 | - mongodb 20 | links: 21 | - mongodb:mongodb 22 | - redis:redis 23 | networks: 24 | - internal 25 | - web 26 | 27 | mongodb: 28 | image: mongo:latest 29 | container_name: imsearch-mongodb 30 | volumes: 31 | - mongo-data:/var/lib/mongodb/db 32 | - mongo-backup:/var/lib/backup 33 | ports: 34 | - 27017:27017 35 | networks: 36 | - internal 37 | command: mongod --replSet mongodb0 --smallfiles 38 | 39 | redis: 40 | image: redis:5-alpine 41 | container_name: imsearch-redis 42 | restart: unless-stopped 43 | ports: 44 | - "6379:6379" 45 | networks: 46 | - internal 47 | 48 | networks: 49 | internal: 50 | web: 51 | external: true 52 | 53 | volumes: 54 | mongo-data: 55 | mongo-backup: 56 | app_data: 57 | driver_opts: 58 | type: none 59 | o: bind 60 | device: ${PWD}/shared/imsearch 61 | -------------------------------------------------------------------------------- /docker-compose.yml: -------------------------------------------------------------------------------- 1 | --- 2 | version: '3.7' 3 | services: 4 | 5 | imsearch: 6 | image: imsearch-cpu:latest 7 | container_name: imsearch-cpu 8 | build: 9 | context: . 10 | dockerfile: Dockerfile 11 | ports: 12 | - 5000:5000 13 | environment: 14 | - MONGO_URI=mongodb://mongodb:27017/ 15 | - REDIS_URI=redis://redis:6379/0 16 | - STORAGE_MODE=local 17 | depends_on: 18 | - redis 19 | - mongodb 20 | links: 21 | - mongodb:mongodb 22 | - redis:redis 23 | networks: 24 | - internal 25 | - web 26 | 27 | mongodb: 28 | image: mongo:latest 29 | container_name: imsearch-mongodb 30 | volumes: 31 | - mongo-data:/var/lib/mongodb/db 32 | - mongo-backup:/var/lib/backup 33 | ports: 34 | - 27017:27017 35 | networks: 36 | - internal 37 | command: mongod --replSet mongodb0 --smallfiles 38 | 39 | redis: 40 | image: redis:5-alpine 41 | container_name: imsearch-redis 42 | restart: unless-stopped 43 | ports: 44 | - "6379:6379" 45 | networks: 46 | - internal 47 | 48 | networks: 49 | internal: 50 | web: 51 | external: true 52 | 53 | volumes: 54 | mongo-data: 55 | mongo-backup: 56 | app_data: 57 | driver_opts: 58 | type: none 59 | o: bind 60 | device: ${PWD}/shared/imsearch 61 | -------------------------------------------------------------------------------- /examples/index.py: -------------------------------------------------------------------------------- 1 | import cv2 2 | import glob 3 | import os 4 | import sys 5 | 6 | import imsearch 7 | 8 | 9 | def create_index(name, images): 10 | # Initialize the index 11 | index = imsearch.init(name) 12 | 13 | # Clear the index and data if any exists 14 | index.cleanIndex() 15 | 16 | # Add single image to index (image path locally stored) 17 | index.addImage(images[0]) 18 | 19 | # Add image using URL with same interface 20 | index.addImage( 21 | "https://www.wallpaperup.com/uploads/wallpapers/2014/04/14/332423/d5c09641cb3af3a18087937d55125ae3-700.jpg") 22 | 23 | # Add images in batch (List of image paths locally stored) 24 | index.addImageBatch(images[1:]) 25 | 26 | # Build the index 27 | index.createIndex() 28 | 29 | return index 30 | 31 | 32 | def create_index_with_config(name): 33 | ''' 34 | parameters: 35 | name: name of the index (unique identifier name) 36 | MONGO_URI 37 | REDIS_URI 38 | DETECTOR_MODE: select 'local' or 'remote'. 39 | local: detector backend should be running on the same machine 40 | remote: Process should not start detector backend 41 | pass any configuration you want to expose as environment variable. 42 | ''' 43 | index = imsearch.init(name=name, 44 | MONGO_URI='mongodb://localhost:27017/', 45 | REDIS_URI='redis://dummy:Welcome00@123.456.78.111:6379/0', 46 | DETECTOR_MODE='local') 47 | 48 | 49 | def show_results(similar, qImage): 50 | qImage = imsearch.utils.check_load_image(qImage) 51 | qImage = cv2.cvtColor(qImage, cv2.COLOR_RGB2BGR) 52 | cv2.imshow('qImage', qImage) 53 | for _i, _s in similar: 54 | rImage = cv2.imread(_i['image']) 55 | print([x['name'] for x in _i['primary']]) 56 | print(_s) 57 | cv2.imshow('rImage', rImage) 58 | cv2.waitKey(0) 59 | 60 | 61 | if __name__ == "__main__": 62 | all_images = glob.glob(os.path.join( 63 | os.path.dirname(__file__), '..', 'images/*.jpg')) 64 | index = create_index('test', all_images) 65 | 66 | # query index with image path 67 | ''' 68 | image_path: path to image or URL 69 | k: Number of results 70 | policy: choose policy from 'object' or 'global'. Search results will change accordingly. 71 | ''' 72 | similar, _ = index.knnQuery(image_path=all_images[-1], k=10, policy='object') 73 | show_results(similar, all_images[-1]) 74 | 75 | # query with image URL 76 | img_url = 'https://www.wallpaperup.com/uploads/wallpapers/2014/04/14/332423/d5c09641cb3af3a18087937d55125ae3-700.jpg' 77 | similar, _ = index.knnQuery(image_path=img_url, k=10, policy='global') 78 | show_results(similar, img_url) 79 | 80 | # Create index with configuration 81 | index = create_index_with_config('test') 82 | -------------------------------------------------------------------------------- /examples/save_index.py: -------------------------------------------------------------------------------- 1 | import cv2 2 | import glob 3 | import os 4 | import sys 5 | 6 | import imsearch 7 | 8 | 9 | def create_and_save(name, images, file_path): 10 | index = imsearch.init(name) 11 | index.cleanIndex() 12 | index.addImageBatch(images) 13 | index.createIndex() 14 | index.saveIndex(file_path) 15 | index.cleanIndex() 16 | 17 | 18 | def load_index(name, file_path): 19 | index = imsearch.init_from_file(file_path=file_path, name=name) 20 | index.createIndex() 21 | return index 22 | 23 | 24 | if __name__ == "__main__": 25 | all_images = glob.glob(os.path.join( 26 | os.path.dirname(__file__), '..', 'images/*.jpg')) 27 | 28 | create_and_save('test', all_images, 'test_index.tar.gz') 29 | index = load_index('test_name', 'test_index.tar.gz') -------------------------------------------------------------------------------- /examples/storage.py: -------------------------------------------------------------------------------- 1 | import cv2 2 | import glob 3 | import os 4 | import sys 5 | import json 6 | 7 | import imsearch 8 | 9 | gcp_config = { 10 | 'GOOGLE_APPLICATION_CREDENTIALS': '../.config/cloud-ml-f1954f23eaa8.json', 11 | 'BUCKET_NAME': 'imsearch-testing', 12 | 'STORAGE_MODE': 'gcp' 13 | } 14 | 15 | with open('../.config/aws-config.json', 'r') as fp: 16 | aws_config_file = json.load(fp) 17 | 18 | aws_config = { 19 | 'AWS_ACCESS_KEY_ID': aws_config_file['AWS_ACCESS_KEY_ID'], 20 | 'AWS_SECRET_ACCESS_KEY': aws_config_file['AWS_SECRET_ACCESS_KEY'], 21 | 'BUCKET_NAME': aws_config_file['BUCKET_NAME'], 22 | 'STORAGE_MODE': 's3' 23 | } 24 | 25 | 26 | def show_results(similar, qImage): 27 | qImage = imsearch.utils.check_load_image(qImage) 28 | qImage = cv2.cvtColor(qImage, cv2.COLOR_RGB2BGR) 29 | cv2.imshow('qImage', qImage) 30 | for _i, _s in similar: 31 | rImage = cv2.cvtColor(imsearch.utils.check_load_image( 32 | _i['image']), cv2.COLOR_RGB2BGR) 33 | print([x['name'] for x in _i['primary']]) 34 | print(_s) 35 | cv2.imshow('rImage', rImage) 36 | cv2.waitKey(0) 37 | 38 | 39 | if __name__ == "__main__": 40 | all_images = glob.glob(os.path.join( 41 | os.path.dirname(__file__), '..', 'images/*.jpg')) 42 | index = imsearch.init(name='test', **aws_config) 43 | index.cleanIndex() 44 | index.addImageBatch(all_images) 45 | index.createIndex() 46 | 47 | # query with image URL 48 | img_url = 'https://www.wallpaperup.com/uploads/wallpapers/2014/04/14/332423/d5c09641cb3af3a18087937d55125ae3-700.jpg' 49 | similar, _ = index.knnQuery(image_path=img_url, k=10, policy='global') 50 | show_results(similar, img_url) 51 | -------------------------------------------------------------------------------- /examples/yolo_detector.py: -------------------------------------------------------------------------------- 1 | from imsearch.backend.object_detector.yolo import Detector 2 | from imsearch import utils 3 | from PIL import Image 4 | import numpy as np 5 | 6 | import matplotlib.pyplot as plt 7 | import matplotlib.patches as patches 8 | from matplotlib.ticker import NullLocator 9 | 10 | 11 | def show_output(output, img): 12 | fig, ax = plt.subplots(1) 13 | ax.imshow(img) 14 | 15 | for obj in output: 16 | box = obj['box'] 17 | x1, y1, x2, y2 = box[0], box[1], box[2], box[3] 18 | box_w = x2 - x1 19 | box_h = y2 - y1 20 | bbox = patches.Rectangle((x1, y1), box_w, box_h, 21 | linewidth=2, edgecolor='red', facecolor="none") 22 | ax.add_patch(bbox) 23 | plt.text( 24 | x1, 25 | y1, 26 | s=obj['name'], 27 | color="white", 28 | verticalalignment="top", 29 | bbox={"color": 'red', "pad": 0}, 30 | ) 31 | 32 | plt.axis("off") 33 | plt.gca().xaxis.set_major_locator(NullLocator()) 34 | plt.gca().yaxis.set_major_locator(NullLocator()) 35 | plt.show() 36 | 37 | 38 | if __name__ == "__main__": 39 | PATH = '../images/000000055167.jpg' 40 | detector = Detector() 41 | 42 | img = utils.check_load_image(PATH) 43 | output = detector.predict(img) 44 | show_output(output, img) 45 | -------------------------------------------------------------------------------- /images/000000055167.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/rikenmehta03/imsearch/0f86ea4c0017b6759e1d79281ba288c5d2641870/images/000000055167.jpg -------------------------------------------------------------------------------- /images/000000084362.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/rikenmehta03/imsearch/0f86ea4c0017b6759e1d79281ba288c5d2641870/images/000000084362.jpg -------------------------------------------------------------------------------- /images/000000101068.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/rikenmehta03/imsearch/0f86ea4c0017b6759e1d79281ba288c5d2641870/images/000000101068.jpg -------------------------------------------------------------------------------- /images/000000134112.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/rikenmehta03/imsearch/0f86ea4c0017b6759e1d79281ba288c5d2641870/images/000000134112.jpg -------------------------------------------------------------------------------- /images/000000157138.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/rikenmehta03/imsearch/0f86ea4c0017b6759e1d79281ba288c5d2641870/images/000000157138.jpg -------------------------------------------------------------------------------- /images/000000170545.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/rikenmehta03/imsearch/0f86ea4c0017b6759e1d79281ba288c5d2641870/images/000000170545.jpg -------------------------------------------------------------------------------- /images/000000263463.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/rikenmehta03/imsearch/0f86ea4c0017b6759e1d79281ba288c5d2641870/images/000000263463.jpg -------------------------------------------------------------------------------- /images/000000275727.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/rikenmehta03/imsearch/0f86ea4c0017b6759e1d79281ba288c5d2641870/images/000000275727.jpg -------------------------------------------------------------------------------- /images/000000304812.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/rikenmehta03/imsearch/0f86ea4c0017b6759e1d79281ba288c5d2641870/images/000000304812.jpg -------------------------------------------------------------------------------- /images/000000306582.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/rikenmehta03/imsearch/0f86ea4c0017b6759e1d79281ba288c5d2641870/images/000000306582.jpg -------------------------------------------------------------------------------- /images/000000327617.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/rikenmehta03/imsearch/0f86ea4c0017b6759e1d79281ba288c5d2641870/images/000000327617.jpg -------------------------------------------------------------------------------- /images/000000404805.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/rikenmehta03/imsearch/0f86ea4c0017b6759e1d79281ba288c5d2641870/images/000000404805.jpg -------------------------------------------------------------------------------- /images/000000468124.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/rikenmehta03/imsearch/0f86ea4c0017b6759e1d79281ba288c5d2641870/images/000000468124.jpg -------------------------------------------------------------------------------- /images/000000482735.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/rikenmehta03/imsearch/0f86ea4c0017b6759e1d79281ba288c5d2641870/images/000000482735.jpg -------------------------------------------------------------------------------- /images/000000534041.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/rikenmehta03/imsearch/0f86ea4c0017b6759e1d79281ba288c5d2641870/images/000000534041.jpg -------------------------------------------------------------------------------- /imsearch/__init__.py: -------------------------------------------------------------------------------- 1 | from .config import config_init 2 | from .index import Index 3 | from .backend import run 4 | 5 | 6 | def init(name='index', **config): 7 | config_init(config) 8 | return Index(name=name) 9 | 10 | 11 | def init_from_file(file_path, name=None, **config): 12 | config_init(config) 13 | return Index.loadFromFile(file_path, name=name) 14 | 15 | 16 | def run_detector(redis_url): 17 | run(redis_url) 18 | -------------------------------------------------------------------------------- /imsearch/backend/__init__.py: -------------------------------------------------------------------------------- 1 | import os 2 | import subprocess 3 | 4 | dir_path = os.path.dirname(os.path.realpath(__file__)) 5 | 6 | 7 | def run(redis_url=None): 8 | env = os.environ.copy() 9 | if redis_url is not None: 10 | env['REDIS_URI'] = redis_url 11 | return subprocess.Popen(['python', os.path.join(dir_path, 'extractor.py')], env=env) 12 | -------------------------------------------------------------------------------- /imsearch/backend/extractor.py: -------------------------------------------------------------------------------- 1 | import os 2 | import sys 3 | import time 4 | import json 5 | import base64 6 | import numpy as np 7 | import redis 8 | 9 | from object_detector import get_detector 10 | from feature_extractor import extract_features 11 | from imsearch import utils 12 | 13 | REDIS_DB = redis.StrictRedis.from_url( 14 | os.environ.get('REDIS_URI', 'redis://localhost:6379/0')) 15 | REDIS_QUEUE = 'image_queue' 16 | BATCH_SIZE = 1 17 | 18 | 19 | def main(): 20 | detector = get_detector('yolo') 21 | while True: 22 | time.sleep(0.01) 23 | _queue = REDIS_DB.lrange(REDIS_QUEUE, 0, BATCH_SIZE - 1) 24 | for _q in _queue: 25 | all_features = { 26 | 'primary': [], 27 | 'object_bitmap': [0 for _ in range(len(detector.classes))] 28 | } 29 | _q = json.loads(_q.decode("utf-8")) 30 | img = utils.base64_decode(_q['image'], _q['shape']) 31 | 32 | all_features['secondary'] = extract_features(img.copy()) 33 | 34 | response = detector.predict(img) 35 | for obj in response: 36 | box = obj['box'] 37 | x1, y1, x2, y2 = box[0], box[1], box[2], box[3] 38 | if(x2-x1 >= 75 and y2-y1 >= 75): 39 | features = extract_features(img[y1:y2, x1:x2]) 40 | all_features['primary'].append({ 41 | 'features': features, 42 | 'label': obj['label'], 43 | 'name': obj['name'], 44 | 'box': obj['box'] 45 | }) 46 | all_features['object_bitmap'][obj['label']] = 1 47 | 48 | REDIS_DB.set(_q['id'], json.dumps(all_features)) 49 | REDIS_DB.ltrim(REDIS_QUEUE, len(_queue), -1) 50 | 51 | 52 | if __name__ == "__main__": 53 | print("Running extractor") 54 | main() 55 | -------------------------------------------------------------------------------- /imsearch/backend/feature_extractor.py: -------------------------------------------------------------------------------- 1 | from PIL import Image 2 | import base64 3 | import torch 4 | import torch.nn as nn 5 | from torchvision import models 6 | import torchvision.transforms as transforms 7 | 8 | 9 | def get_extractor(arch='resnet50'): 10 | model_ft = models.__dict__[arch](pretrained=True) 11 | extractor = nn.Sequential(*list(model_ft.children())[:-1]) 12 | return extractor 13 | 14 | 15 | device = torch.device("cuda" if torch.cuda.is_available() else "cpu") 16 | extractor = get_extractor('resnet50').to(device) 17 | extractor.eval() 18 | image_transforms = transforms.Compose([ 19 | transforms.Resize(300), 20 | transforms.CenterCrop(244), 21 | transforms.ToTensor(), 22 | ]) 23 | 24 | 25 | def extract_features(img): 26 | img = image_transforms(Image.fromarray(img)) 27 | img = img.unsqueeze_(0).to(device) 28 | 29 | with torch.no_grad(): 30 | out = extractor(img).squeeze().detach().cpu() 31 | return base64.b64encode(out.numpy()).decode("utf-8") 32 | -------------------------------------------------------------------------------- /imsearch/backend/object_detector/__init__.py: -------------------------------------------------------------------------------- 1 | from .yolo import Detector 2 | 3 | 4 | def get_detector(detector_type='yolo'): 5 | if detector_type == 'yolo': 6 | return Detector() 7 | -------------------------------------------------------------------------------- /imsearch/backend/object_detector/yolo/__init__.py: -------------------------------------------------------------------------------- 1 | from .detector import Detector 2 | -------------------------------------------------------------------------------- /imsearch/backend/object_detector/yolo/detector.py: -------------------------------------------------------------------------------- 1 | from __future__ import division 2 | 3 | import os 4 | import sys 5 | import torch 6 | import torchvision.transforms as transforms 7 | from PIL import Image 8 | 9 | from .models import Darknet 10 | from .utils import load_classes, pad_to_square, resize, non_max_suppression, rescale_boxes, download_files 11 | 12 | device = torch.device("cuda" if torch.cuda.is_available() else "cpu") 13 | DATA_PATH = os.path.join(os.environ.get('HOME'), '.imsearch', 'yolo') 14 | 15 | 16 | class Detector(object): 17 | model = None 18 | 19 | def __init__(self): 20 | download_files() 21 | 22 | self.img_size = 416 23 | self.config_path = os.path.join(DATA_PATH, 'yolov3.cfg') 24 | self.weights_path = os.path.join(DATA_PATH, 'yolov3.weights') 25 | self.class_path = os.path.join(DATA_PATH, 'coco.names') 26 | self.conf_thres = 0.8 27 | self.nms_thres = 0.4 28 | if Detector.model is None: 29 | model = Darknet(self.config_path, 30 | img_size=self.img_size).to(device) 31 | model.load_darknet_weights(self.weights_path) 32 | model.eval() 33 | Detector.model = model 34 | self.classes = load_classes(self.class_path) 35 | 36 | def _load_image(self, img): 37 | actual_size = img.shape 38 | img = self._process_image(img) 39 | img = img.unsqueeze_(0) 40 | return img, actual_size 41 | 42 | def _process_image(self, img): 43 | img = Image.fromarray(img) 44 | img = transforms.ToTensor()(img) 45 | img, _ = pad_to_square(img, 0) 46 | img = resize(img, self.img_size) 47 | return img 48 | 49 | def predict(self, img): 50 | input_img, actual_shape = self._load_image(img) 51 | input_img = input_img.to(device) 52 | 53 | with torch.no_grad(): 54 | detections = Detector.model(input_img) 55 | detections = non_max_suppression( 56 | detections, self.conf_thres, self.nms_thres)[0] 57 | 58 | response = [] 59 | if detections is not None: 60 | detections = rescale_boxes( 61 | detections, self.img_size, actual_shape[:2]) 62 | for x1, y1, x2, y2, _, cls_conf, cls_pred in detections: 63 | x1 = max(min(int(x1.item()), actual_shape[1]-1), 0) 64 | x2 = max(min(int(x2.item()), actual_shape[1]-1), 0) 65 | y1 = max(min(int(y1.item()), actual_shape[0]-1), 0) 66 | y2 = max(min(int(y2.item()), actual_shape[0]-1), 0) 67 | response.append({ 68 | 'box': [x1, y1, x2, y2], 69 | 'score': cls_conf.item(), 70 | 'label': int(cls_pred), 71 | 'name': self.classes[int(cls_pred)] 72 | }) 73 | return response 74 | -------------------------------------------------------------------------------- /imsearch/backend/object_detector/yolo/models.py: -------------------------------------------------------------------------------- 1 | from __future__ import division 2 | 3 | import torch 4 | import torch.nn as nn 5 | import torch.nn.functional as F 6 | import numpy as np 7 | 8 | from .utils import build_targets, to_cpu, parse_model_config 9 | 10 | 11 | def create_modules(module_defs): 12 | """ 13 | Constructs module list of layer blocks from module configuration in module_defs 14 | """ 15 | hyperparams = module_defs.pop(0) 16 | output_filters = [int(hyperparams["channels"])] 17 | module_list = nn.ModuleList() 18 | for module_i, module_def in enumerate(module_defs): 19 | modules = nn.Sequential() 20 | 21 | if module_def["type"] == "convolutional": 22 | bn = int(module_def["batch_normalize"]) 23 | filters = int(module_def["filters"]) 24 | kernel_size = int(module_def["size"]) 25 | pad = (kernel_size - 1) // 2 26 | modules.add_module( 27 | f"conv_{module_i}", 28 | nn.Conv2d( 29 | in_channels=output_filters[-1], 30 | out_channels=filters, 31 | kernel_size=kernel_size, 32 | stride=int(module_def["stride"]), 33 | padding=pad, 34 | bias=not bn, 35 | ), 36 | ) 37 | if bn: 38 | modules.add_module(f"batch_norm_{module_i}", nn.BatchNorm2d( 39 | filters, momentum=0.9, eps=1e-5)) 40 | if module_def["activation"] == "leaky": 41 | modules.add_module(f"leaky_{module_i}", nn.LeakyReLU(0.1)) 42 | 43 | elif module_def["type"] == "maxpool": 44 | kernel_size = int(module_def["size"]) 45 | stride = int(module_def["stride"]) 46 | if kernel_size == 2 and stride == 1: 47 | modules.add_module( 48 | f"_debug_padding_{module_i}", nn.ZeroPad2d((0, 1, 0, 1))) 49 | maxpool = nn.MaxPool2d( 50 | kernel_size=kernel_size, stride=stride, padding=int((kernel_size - 1) // 2)) 51 | modules.add_module(f"maxpool_{module_i}", maxpool) 52 | 53 | elif module_def["type"] == "upsample": 54 | upsample = Upsample(scale_factor=int( 55 | module_def["stride"]), mode="nearest") 56 | modules.add_module(f"upsample_{module_i}", upsample) 57 | 58 | elif module_def["type"] == "route": 59 | layers = [int(x) for x in module_def["layers"].split(",")] 60 | filters = sum([output_filters[1:][i] for i in layers]) 61 | modules.add_module(f"route_{module_i}", EmptyLayer()) 62 | 63 | elif module_def["type"] == "shortcut": 64 | filters = output_filters[1:][int(module_def["from"])] 65 | modules.add_module(f"shortcut_{module_i}", EmptyLayer()) 66 | 67 | elif module_def["type"] == "yolo": 68 | anchor_idxs = [int(x) for x in module_def["mask"].split(",")] 69 | # Extract anchors 70 | anchors = [int(x) for x in module_def["anchors"].split(",")] 71 | anchors = [(anchors[i], anchors[i + 1]) 72 | for i in range(0, len(anchors), 2)] 73 | anchors = [anchors[i] for i in anchor_idxs] 74 | num_classes = int(module_def["classes"]) 75 | img_size = int(hyperparams["height"]) 76 | # Define detection layer 77 | yolo_layer = YOLOLayer(anchors, num_classes, img_size) 78 | modules.add_module(f"yolo_{module_i}", yolo_layer) 79 | # Register module list and number of output filters 80 | module_list.append(modules) 81 | output_filters.append(filters) 82 | 83 | return hyperparams, module_list 84 | 85 | 86 | class Upsample(nn.Module): 87 | """ nn.Upsample is deprecated """ 88 | 89 | def __init__(self, scale_factor, mode="nearest"): 90 | super(Upsample, self).__init__() 91 | self.scale_factor = scale_factor 92 | self.mode = mode 93 | 94 | def forward(self, x): 95 | x = F.interpolate(x, scale_factor=self.scale_factor, mode=self.mode) 96 | return x 97 | 98 | 99 | class EmptyLayer(nn.Module): 100 | """Placeholder for 'route' and 'shortcut' layers""" 101 | 102 | def __init__(self): 103 | super(EmptyLayer, self).__init__() 104 | 105 | 106 | class YOLOLayer(nn.Module): 107 | """Detection layer""" 108 | 109 | def __init__(self, anchors, num_classes, img_dim=416): 110 | super(YOLOLayer, self).__init__() 111 | self.anchors = anchors 112 | self.num_anchors = len(anchors) 113 | self.num_classes = num_classes 114 | self.ignore_thres = 0.5 115 | self.mse_loss = nn.MSELoss() 116 | self.bce_loss = nn.BCELoss() 117 | self.obj_scale = 1 118 | self.noobj_scale = 100 119 | self.metrics = {} 120 | self.img_dim = img_dim 121 | self.grid_size = 0 # grid size 122 | 123 | def compute_grid_offsets(self, grid_size, cuda=True): 124 | self.grid_size = grid_size 125 | g = self.grid_size 126 | FloatTensor = torch.cuda.FloatTensor if cuda else torch.FloatTensor 127 | self.stride = self.img_dim / self.grid_size 128 | # Calculate offsets for each grid 129 | self.grid_x = torch.arange(g).repeat( 130 | g, 1).view([1, 1, g, g]).type(FloatTensor) 131 | self.grid_y = torch.arange(g).repeat( 132 | g, 1).t().view([1, 1, g, g]).type(FloatTensor) 133 | self.scaled_anchors = FloatTensor( 134 | [(a_w / self.stride, a_h / self.stride) for a_w, a_h in self.anchors]) 135 | self.anchor_w = self.scaled_anchors[:, 0:1].view( 136 | (1, self.num_anchors, 1, 1)) 137 | self.anchor_h = self.scaled_anchors[:, 1:2].view( 138 | (1, self.num_anchors, 1, 1)) 139 | 140 | def forward(self, x, targets=None, img_dim=None): 141 | 142 | # Tensors for cuda support 143 | FloatTensor = torch.cuda.FloatTensor if x.is_cuda else torch.FloatTensor 144 | LongTensor = torch.cuda.LongTensor if x.is_cuda else torch.LongTensor 145 | ByteTensor = torch.cuda.ByteTensor if x.is_cuda else torch.ByteTensor 146 | 147 | self.img_dim = img_dim 148 | num_samples = x.size(0) 149 | grid_size = x.size(2) 150 | 151 | prediction = ( 152 | x.view(num_samples, self.num_anchors, 153 | self.num_classes + 5, grid_size, grid_size) 154 | .permute(0, 1, 3, 4, 2) 155 | .contiguous() 156 | ) 157 | 158 | # Get outputs 159 | x = torch.sigmoid(prediction[..., 0]) # Center x 160 | y = torch.sigmoid(prediction[..., 1]) # Center y 161 | w = prediction[..., 2] # Width 162 | h = prediction[..., 3] # Height 163 | pred_conf = torch.sigmoid(prediction[..., 4]) # Conf 164 | pred_cls = torch.sigmoid(prediction[..., 5:]) # Cls pred. 165 | 166 | # If grid size does not match current we compute new offsets 167 | if grid_size != self.grid_size: 168 | self.compute_grid_offsets(grid_size, cuda=x.is_cuda) 169 | 170 | # Add offset and scale with anchors 171 | pred_boxes = FloatTensor(prediction[..., :4].shape) 172 | pred_boxes[..., 0] = x.data + self.grid_x 173 | pred_boxes[..., 1] = y.data + self.grid_y 174 | pred_boxes[..., 2] = torch.exp(w.data) * self.anchor_w 175 | pred_boxes[..., 3] = torch.exp(h.data) * self.anchor_h 176 | 177 | output = torch.cat( 178 | ( 179 | pred_boxes.view(num_samples, -1, 4) * self.stride, 180 | pred_conf.view(num_samples, -1, 1), 181 | pred_cls.view(num_samples, -1, self.num_classes), 182 | ), 183 | -1, 184 | ) 185 | 186 | if targets is None: 187 | return output, 0 188 | else: 189 | iou_scores, class_mask, obj_mask, noobj_mask, tx, ty, tw, th, tcls, tconf = build_targets( 190 | pred_boxes=pred_boxes, 191 | pred_cls=pred_cls, 192 | target=targets, 193 | anchors=self.scaled_anchors, 194 | ignore_thres=self.ignore_thres, 195 | ) 196 | 197 | # Loss : Mask outputs to ignore non-existing objects (except with conf. loss) 198 | loss_x = self.mse_loss(x[obj_mask], tx[obj_mask]) 199 | loss_y = self.mse_loss(y[obj_mask], ty[obj_mask]) 200 | loss_w = self.mse_loss(w[obj_mask], tw[obj_mask]) 201 | loss_h = self.mse_loss(h[obj_mask], th[obj_mask]) 202 | loss_conf_obj = self.bce_loss(pred_conf[obj_mask], tconf[obj_mask]) 203 | loss_conf_noobj = self.bce_loss( 204 | pred_conf[noobj_mask], tconf[noobj_mask]) 205 | loss_conf = self.obj_scale * loss_conf_obj + self.noobj_scale * loss_conf_noobj 206 | loss_cls = self.bce_loss(pred_cls[obj_mask], tcls[obj_mask]) 207 | total_loss = loss_x + loss_y + loss_w + loss_h + loss_conf + loss_cls 208 | 209 | # Metrics 210 | cls_acc = 100 * class_mask[obj_mask].mean() 211 | conf_obj = pred_conf[obj_mask].mean() 212 | conf_noobj = pred_conf[noobj_mask].mean() 213 | conf50 = (pred_conf > 0.5).float() 214 | iou50 = (iou_scores > 0.5).float() 215 | iou75 = (iou_scores > 0.75).float() 216 | detected_mask = conf50 * class_mask * tconf 217 | precision = torch.sum(iou50 * detected_mask) / \ 218 | (conf50.sum() + 1e-16) 219 | recall50 = torch.sum(iou50 * detected_mask) / \ 220 | (obj_mask.sum() + 1e-16) 221 | recall75 = torch.sum(iou75 * detected_mask) / \ 222 | (obj_mask.sum() + 1e-16) 223 | 224 | self.metrics = { 225 | "loss": to_cpu(total_loss).item(), 226 | "x": to_cpu(loss_x).item(), 227 | "y": to_cpu(loss_y).item(), 228 | "w": to_cpu(loss_w).item(), 229 | "h": to_cpu(loss_h).item(), 230 | "conf": to_cpu(loss_conf).item(), 231 | "cls": to_cpu(loss_cls).item(), 232 | "cls_acc": to_cpu(cls_acc).item(), 233 | "recall50": to_cpu(recall50).item(), 234 | "recall75": to_cpu(recall75).item(), 235 | "precision": to_cpu(precision).item(), 236 | "conf_obj": to_cpu(conf_obj).item(), 237 | "conf_noobj": to_cpu(conf_noobj).item(), 238 | "grid_size": grid_size, 239 | } 240 | 241 | return output, total_loss 242 | 243 | 244 | class Darknet(nn.Module): 245 | """YOLOv3 object detection model""" 246 | 247 | def __init__(self, config_path, img_size=416): 248 | super(Darknet, self).__init__() 249 | self.module_defs = parse_model_config(config_path) 250 | self.hyperparams, self.module_list = create_modules(self.module_defs) 251 | self.yolo_layers = [layer[0] 252 | for layer in self.module_list if hasattr(layer[0], "metrics")] 253 | self.img_size = img_size 254 | self.seen = 0 255 | self.header_info = np.array([0, 0, 0, self.seen, 0], dtype=np.int32) 256 | 257 | def forward(self, x, targets=None): 258 | img_dim = x.shape[2] 259 | loss = 0 260 | layer_outputs, yolo_outputs = [], [] 261 | for i, (module_def, module) in enumerate(zip(self.module_defs, self.module_list)): 262 | if module_def["type"] in ["convolutional", "upsample", "maxpool"]: 263 | x = module(x) 264 | elif module_def["type"] == "route": 265 | x = torch.cat([layer_outputs[int(layer_i)] 266 | for layer_i in module_def["layers"].split(",")], 1) 267 | elif module_def["type"] == "shortcut": 268 | layer_i = int(module_def["from"]) 269 | x = layer_outputs[-1] + layer_outputs[layer_i] 270 | elif module_def["type"] == "yolo": 271 | x, layer_loss = module[0](x, targets, img_dim) 272 | loss += layer_loss 273 | yolo_outputs.append(x) 274 | layer_outputs.append(x) 275 | yolo_outputs = to_cpu(torch.cat(yolo_outputs, 1)) 276 | return yolo_outputs if targets is None else (loss, yolo_outputs) 277 | 278 | def load_darknet_weights(self, weights_path): 279 | """Parses and loads the weights stored in 'weights_path'""" 280 | 281 | # Open the weights file 282 | with open(weights_path, "rb") as f: 283 | # First five are header values 284 | header = np.fromfile(f, dtype=np.int32, count=5) 285 | self.header_info = header # Needed to write header when saving weights 286 | self.seen = header[3] # number of images seen during training 287 | weights = np.fromfile(f, dtype=np.float32) # The rest are weights 288 | 289 | # Establish cutoff for loading backbone weights 290 | cutoff = None 291 | if "darknet53.conv.74" in weights_path: 292 | cutoff = 75 293 | 294 | ptr = 0 295 | for i, (module_def, module) in enumerate(zip(self.module_defs, self.module_list)): 296 | if i == cutoff: 297 | break 298 | if module_def["type"] == "convolutional": 299 | conv_layer = module[0] 300 | if module_def["batch_normalize"]: 301 | # Load BN bias, weights, running mean and running variance 302 | bn_layer = module[1] 303 | num_b = bn_layer.bias.numel() # Number of biases 304 | # Bias 305 | bn_b = torch.from_numpy( 306 | weights[ptr: ptr + num_b]).view_as(bn_layer.bias) 307 | bn_layer.bias.data.copy_(bn_b) 308 | ptr += num_b 309 | # Weight 310 | bn_w = torch.from_numpy( 311 | weights[ptr: ptr + num_b]).view_as(bn_layer.weight) 312 | bn_layer.weight.data.copy_(bn_w) 313 | ptr += num_b 314 | # Running Mean 315 | bn_rm = torch.from_numpy( 316 | weights[ptr: ptr + num_b]).view_as(bn_layer.running_mean) 317 | bn_layer.running_mean.data.copy_(bn_rm) 318 | ptr += num_b 319 | # Running Var 320 | bn_rv = torch.from_numpy( 321 | weights[ptr: ptr + num_b]).view_as(bn_layer.running_var) 322 | bn_layer.running_var.data.copy_(bn_rv) 323 | ptr += num_b 324 | else: 325 | # Load conv. bias 326 | num_b = conv_layer.bias.numel() 327 | conv_b = torch.from_numpy( 328 | weights[ptr: ptr + num_b]).view_as(conv_layer.bias) 329 | conv_layer.bias.data.copy_(conv_b) 330 | ptr += num_b 331 | # Load conv. weights 332 | num_w = conv_layer.weight.numel() 333 | conv_w = torch.from_numpy( 334 | weights[ptr: ptr + num_w]).view_as(conv_layer.weight) 335 | conv_layer.weight.data.copy_(conv_w) 336 | ptr += num_w 337 | 338 | def save_darknet_weights(self, path, cutoff=-1): 339 | """ 340 | @:param path - path of the new weights file 341 | @:param cutoff - save layers between 0 and cutoff (cutoff = -1 -> all are saved) 342 | """ 343 | fp = open(path, "wb") 344 | self.header_info[3] = self.seen 345 | self.header_info.tofile(fp) 346 | 347 | # Iterate through layers 348 | for i, (module_def, module) in enumerate(zip(self.module_defs[:cutoff], self.module_list[:cutoff])): 349 | if module_def["type"] == "convolutional": 350 | conv_layer = module[0] 351 | # If batch norm, load bn first 352 | if module_def["batch_normalize"]: 353 | bn_layer = module[1] 354 | bn_layer.bias.data.cpu().numpy().tofile(fp) 355 | bn_layer.weight.data.cpu().numpy().tofile(fp) 356 | bn_layer.running_mean.data.cpu().numpy().tofile(fp) 357 | bn_layer.running_var.data.cpu().numpy().tofile(fp) 358 | # Load conv bias 359 | else: 360 | conv_layer.bias.data.cpu().numpy().tofile(fp) 361 | # Load conv weights 362 | conv_layer.weight.data.cpu().numpy().tofile(fp) 363 | 364 | fp.close() 365 | -------------------------------------------------------------------------------- /imsearch/backend/object_detector/yolo/utils.py: -------------------------------------------------------------------------------- 1 | from __future__ import division 2 | import os 3 | import wget 4 | import torch 5 | import torch.nn.functional as F 6 | import numpy as np 7 | 8 | 9 | def to_cpu(tensor): 10 | return tensor.detach().cpu() 11 | 12 | 13 | def load_classes(path): 14 | """ 15 | Loads class labels at 'path' 16 | """ 17 | fp = open(path, "r") 18 | names = fp.read().split("\n")[:-1] 19 | return names 20 | 21 | 22 | def rescale_boxes(boxes, current_dim, original_shape): 23 | """ Rescales bounding boxes to the original shape """ 24 | orig_h, orig_w = original_shape 25 | # The amount of padding that was added 26 | pad_x = max(orig_h - orig_w, 0) * (current_dim / max(original_shape)) 27 | pad_y = max(orig_w - orig_h, 0) * (current_dim / max(original_shape)) 28 | # Image height and width after padding is removed 29 | unpad_h = current_dim - pad_y 30 | unpad_w = current_dim - pad_x 31 | # Rescale bounding boxes to dimension of original image 32 | boxes[:, 0] = ((boxes[:, 0] - pad_x // 2) / unpad_w) * orig_w 33 | boxes[:, 1] = ((boxes[:, 1] - pad_y // 2) / unpad_h) * orig_h 34 | boxes[:, 2] = ((boxes[:, 2] - pad_x // 2) / unpad_w) * orig_w 35 | boxes[:, 3] = ((boxes[:, 3] - pad_y // 2) / unpad_h) * orig_h 36 | return boxes 37 | 38 | 39 | def xywh2xyxy(x): 40 | y = x.new(x.shape) 41 | y[..., 0] = x[..., 0] - x[..., 2] / 2 42 | y[..., 1] = x[..., 1] - x[..., 3] / 2 43 | y[..., 2] = x[..., 0] + x[..., 2] / 2 44 | y[..., 3] = x[..., 1] + x[..., 3] / 2 45 | return y 46 | 47 | 48 | def bbox_wh_iou(wh1, wh2): 49 | wh2 = wh2.t() 50 | w1, h1 = wh1[0], wh1[1] 51 | w2, h2 = wh2[0], wh2[1] 52 | inter_area = torch.min(w1, w2) * torch.min(h1, h2) 53 | union_area = (w1 * h1 + 1e-16) + w2 * h2 - inter_area 54 | return inter_area / union_area 55 | 56 | 57 | def bbox_iou(box1, box2, x1y1x2y2=True): 58 | """ 59 | Returns the IoU of two bounding boxes 60 | """ 61 | if not x1y1x2y2: 62 | # Transform from center and width to exact coordinates 63 | b1_x1, b1_x2 = box1[:, 0] - box1[:, 2] / 2, box1[:, 0] + box1[:, 2] / 2 64 | b1_y1, b1_y2 = box1[:, 1] - box1[:, 3] / 2, box1[:, 1] + box1[:, 3] / 2 65 | b2_x1, b2_x2 = box2[:, 0] - box2[:, 2] / 2, box2[:, 0] + box2[:, 2] / 2 66 | b2_y1, b2_y2 = box2[:, 1] - box2[:, 3] / 2, box2[:, 1] + box2[:, 3] / 2 67 | else: 68 | # Get the coordinates of bounding boxes 69 | b1_x1, b1_y1, b1_x2, b1_y2 = box1[:, 70 | 0], box1[:, 1], box1[:, 2], box1[:, 3] 71 | b2_x1, b2_y1, b2_x2, b2_y2 = box2[:, 72 | 0], box2[:, 1], box2[:, 2], box2[:, 3] 73 | 74 | # get the corrdinates of the intersection rectangle 75 | inter_rect_x1 = torch.max(b1_x1, b2_x1) 76 | inter_rect_y1 = torch.max(b1_y1, b2_y1) 77 | inter_rect_x2 = torch.min(b1_x2, b2_x2) 78 | inter_rect_y2 = torch.min(b1_y2, b2_y2) 79 | # Intersection area 80 | inter_area = torch.clamp(inter_rect_x2 - inter_rect_x1 + 1, min=0) * torch.clamp( 81 | inter_rect_y2 - inter_rect_y1 + 1, min=0 82 | ) 83 | # Union Area 84 | b1_area = (b1_x2 - b1_x1 + 1) * (b1_y2 - b1_y1 + 1) 85 | b2_area = (b2_x2 - b2_x1 + 1) * (b2_y2 - b2_y1 + 1) 86 | 87 | iou = inter_area / (b1_area + b2_area - inter_area + 1e-16) 88 | 89 | return iou 90 | 91 | 92 | def non_max_suppression(prediction, conf_thres=0.5, nms_thres=0.4): 93 | """ 94 | Removes detections with lower object confidence score than 'conf_thres' and performs 95 | Non-Maximum Suppression to further filter detections. 96 | Returns detections with shape: 97 | (x1, y1, x2, y2, object_conf, class_score, class_pred) 98 | """ 99 | 100 | # From (center x, center y, width, height) to (x1, y1, x2, y2) 101 | prediction[..., :4] = xywh2xyxy(prediction[..., :4]) 102 | output = [None for _ in range(len(prediction))] 103 | for image_i, image_pred in enumerate(prediction): 104 | # Filter out confidence scores below threshold 105 | image_pred = image_pred[image_pred[:, 4] >= conf_thres] 106 | # If none are remaining => process next image 107 | if not image_pred.size(0): 108 | continue 109 | # Object confidence times class confidence 110 | score = image_pred[:, 4] * image_pred[:, 5:].max(1)[0] 111 | # Sort by it 112 | image_pred = image_pred[(-score).argsort()] 113 | class_confs, class_preds = image_pred[:, 5:].max(1, keepdim=True) 114 | detections = torch.cat( 115 | (image_pred[:, :5], class_confs.float(), class_preds.float()), 1) 116 | # Perform non-maximum suppression 117 | keep_boxes = [] 118 | while detections.size(0): 119 | large_overlap = bbox_iou(detections[0, :4].unsqueeze( 120 | 0), detections[:, :4]) > nms_thres 121 | label_match = detections[0, -1] == detections[:, -1] 122 | # Indices of boxes with lower confidence scores, large IOUs and matching labels 123 | invalid = large_overlap & label_match 124 | weights = detections[invalid, 4:5] 125 | # Merge overlapping bboxes by order of confidence 126 | detections[0, :4] = ( 127 | weights * detections[invalid, :4]).sum(0) / weights.sum() 128 | keep_boxes += [detections[0]] 129 | detections = detections[~invalid] 130 | if keep_boxes: 131 | output[image_i] = torch.stack(keep_boxes) 132 | 133 | return output 134 | 135 | 136 | def build_targets(pred_boxes, pred_cls, target, anchors, ignore_thres): 137 | 138 | ByteTensor = torch.cuda.ByteTensor if pred_boxes.is_cuda else torch.ByteTensor 139 | FloatTensor = torch.cuda.FloatTensor if pred_boxes.is_cuda else torch.FloatTensor 140 | 141 | nB = pred_boxes.size(0) 142 | nA = pred_boxes.size(1) 143 | nC = pred_cls.size(-1) 144 | nG = pred_boxes.size(2) 145 | 146 | # Output tensors 147 | obj_mask = ByteTensor(nB, nA, nG, nG).fill_(0) 148 | noobj_mask = ByteTensor(nB, nA, nG, nG).fill_(1) 149 | class_mask = FloatTensor(nB, nA, nG, nG).fill_(0) 150 | iou_scores = FloatTensor(nB, nA, nG, nG).fill_(0) 151 | tx = FloatTensor(nB, nA, nG, nG).fill_(0) 152 | ty = FloatTensor(nB, nA, nG, nG).fill_(0) 153 | tw = FloatTensor(nB, nA, nG, nG).fill_(0) 154 | th = FloatTensor(nB, nA, nG, nG).fill_(0) 155 | tcls = FloatTensor(nB, nA, nG, nG, nC).fill_(0) 156 | 157 | # Convert to position relative to box 158 | target_boxes = target[:, 2:6] * nG 159 | gxy = target_boxes[:, :2] 160 | gwh = target_boxes[:, 2:] 161 | # Get anchors with best iou 162 | ious = torch.stack([bbox_wh_iou(anchor, gwh) for anchor in anchors]) 163 | best_ious, best_n = ious.max(0) 164 | # Separate target values 165 | b, target_labels = target[:, :2].long().t() 166 | gx, gy = gxy.t() 167 | gw, gh = gwh.t() 168 | gi, gj = gxy.long().t() 169 | # Set masks 170 | obj_mask[b, best_n, gj, gi] = 1 171 | noobj_mask[b, best_n, gj, gi] = 0 172 | 173 | # Set noobj mask to zero where iou exceeds ignore threshold 174 | for i, anchor_ious in enumerate(ious.t()): 175 | noobj_mask[b[i], anchor_ious > ignore_thres, gj[i], gi[i]] = 0 176 | 177 | # Coordinates 178 | tx[b, best_n, gj, gi] = gx - gx.floor() 179 | ty[b, best_n, gj, gi] = gy - gy.floor() 180 | # Width and height 181 | tw[b, best_n, gj, gi] = torch.log(gw / anchors[best_n][:, 0] + 1e-16) 182 | th[b, best_n, gj, gi] = torch.log(gh / anchors[best_n][:, 1] + 1e-16) 183 | # One-hot encoding of label 184 | tcls[b, best_n, gj, gi, target_labels] = 1 185 | # Compute label correctness and iou at best anchor 186 | class_mask[b, best_n, gj, gi] = ( 187 | pred_cls[b, best_n, gj, gi].argmax(-1) == target_labels).float() 188 | iou_scores[b, best_n, gj, gi] = bbox_iou( 189 | pred_boxes[b, best_n, gj, gi], target_boxes, x1y1x2y2=False) 190 | 191 | tconf = obj_mask.float() 192 | return iou_scores, class_mask, obj_mask, noobj_mask, tx, ty, tw, th, tcls, tconf 193 | 194 | 195 | def parse_model_config(path): 196 | """Parses the yolo-v3 layer configuration file and returns module definitions""" 197 | file = open(path, 'r') 198 | lines = file.read().split('\n') 199 | lines = [x for x in lines if x and not x.startswith('#')] 200 | lines = [x.rstrip().lstrip() 201 | for x in lines] # get rid of fringe whitespaces 202 | module_defs = [] 203 | for line in lines: 204 | if line.startswith('['): # This marks the start of a new block 205 | module_defs.append({}) 206 | module_defs[-1]['type'] = line[1:-1].rstrip() 207 | if module_defs[-1]['type'] == 'convolutional': 208 | module_defs[-1]['batch_normalize'] = 0 209 | else: 210 | key, value = line.split("=") 211 | value = value.strip() 212 | module_defs[-1][key.rstrip()] = value.strip() 213 | 214 | return module_defs 215 | 216 | 217 | def pad_to_square(img, pad_value): 218 | c, h, w = img.shape 219 | dim_diff = np.abs(h - w) 220 | # (upper / left) padding and (lower / right) padding 221 | pad1, pad2 = dim_diff // 2, dim_diff - dim_diff // 2 222 | # Determine padding 223 | pad = (0, 0, pad1, pad2) if h <= w else (pad1, pad2, 0, 0) 224 | # Add padding 225 | img = F.pad(img, pad, "constant", value=pad_value) 226 | 227 | return img, pad 228 | 229 | 230 | def resize(image, size): 231 | image = F.interpolate(image.unsqueeze(0), size=size, 232 | mode="nearest").squeeze(0) 233 | return image 234 | 235 | 236 | def download_files(): 237 | path = os.path.join(os.environ.get('HOME'), '.imsearch', 'yolo') 238 | os.makedirs(path, exist_ok=True) 239 | config_path = os.path.join(path, 'yolov3.cfg') 240 | if not os.path.exists(config_path): 241 | wget.download( 242 | 'https://www.dropbox.com/s/9kacdlcdo717i1g/yolov3.cfg?dl=1', out=config_path) 243 | 244 | weights_path = os.path.join(path, 'yolov3.weights') 245 | if not os.path.exists(weights_path): 246 | wget.download( 247 | 'https://www.dropbox.com/s/8ftd60s7z77ltv4/yolov3.weights?dl=1', out=weights_path) 248 | 249 | class_path = os.path.join(path, 'coco.names') 250 | if not os.path.exists(class_path): 251 | wget.download( 252 | 'https://www.dropbox.com/s/96o871of6g0secz/coco.names?dl=1', out=class_path) 253 | -------------------------------------------------------------------------------- /imsearch/config.py: -------------------------------------------------------------------------------- 1 | import os 2 | import copy 3 | 4 | default_config = { 5 | 'MONGO_URI': 'mongodb://localhost:27017/', 6 | 'REDIS_URI': 'redis://localhost:6379/0', 7 | 'DETECTOR_MODE': 'local', 8 | 'STORAGE_MODE': 'local' 9 | } 10 | 11 | 12 | def config_init(config): 13 | final_config = {} 14 | for k, v in default_config.items(): 15 | final_config[k] = os.environ.get(k, v) 16 | 17 | for k, v in config.items(): 18 | final_config[k] = v 19 | os.environ.update(final_config) 20 | -------------------------------------------------------------------------------- /imsearch/exception.py: -------------------------------------------------------------------------------- 1 | class InvalidAttributeError(Exception): 2 | """Exception raised for errors in the input. 3 | 4 | Attributes: 5 | expression -- input expression in which the error occurred 6 | message -- explanation of the error 7 | """ 8 | 9 | def __init__(self, expression, message): 10 | self.expression = expression 11 | self.message = message 12 | -------------------------------------------------------------------------------- /imsearch/extractor.py: -------------------------------------------------------------------------------- 1 | import os 2 | import sys 3 | import shutil 4 | import base64 5 | import uuid 6 | import json 7 | import time 8 | from tempfile import NamedTemporaryFile 9 | from PIL import Image 10 | import numpy as np 11 | import redis 12 | 13 | from . import utils 14 | from .storage import get_storage_object 15 | 16 | 17 | class FeatureExtractor: 18 | def __init__(self, index_name): 19 | self.index_name = index_name 20 | self.redis_url = os.environ.get('REDIS_URI') 21 | self._redis_db = redis.StrictRedis.from_url(self.redis_url) 22 | if os.environ.get('STORAGE_MODE') not in ['local', 'none']: 23 | self._storage = get_storage_object(os.environ.get('STORAGE_MODE')) 24 | else: 25 | self._storage = os.environ.get('STORAGE_MODE') 26 | os.makedirs(self._get_image_path(self.index_name), exist_ok=True) 27 | 28 | @classmethod 29 | def _get_image_path(self, index_name): 30 | home_dir = os.environ.get('HOME') 31 | return os.path.join(home_dir, '.imsearch', 'images', index_name) 32 | 33 | def _save_image(self, image, _id): 34 | if self._storage == 'none': 35 | return '' 36 | elif self._storage == 'local': 37 | dst = os.path.join(self._get_image_path( 38 | self.index_name), '{}.jpg'.format(_id)) 39 | utils.save_image(image, dst) 40 | return dst 41 | else: 42 | with NamedTemporaryFile() as temp: 43 | dst = "{}.jpg".format(temp.name) 44 | key = "images/{}/{}.jpg".format(self.index_name, _id) 45 | utils.save_image(image, dst) 46 | return self._storage.upload(dst, key) 47 | 48 | def _decode_redis_data(self, data): 49 | data = json.loads(data.decode('utf-8')) 50 | 51 | def _decode(d): 52 | d['features'] = utils.base64_decode( 53 | d['features'], dtype=np.float32) 54 | return d 55 | 56 | data['primary'] = list(map(_decode, data['primary'])) 57 | data['secondary'] = utils.base64_decode( 58 | data['secondary'], dtype=np.float32) 59 | return data 60 | 61 | def clean(self): 62 | shutil.rmtree(self._get_image_path(self.index_name)) 63 | os.makedirs(self._get_image_path(self.index_name), exist_ok=True) 64 | 65 | def extract(self, image_path, save=True): 66 | image = utils.check_load_image(image_path) 67 | if image is None: 68 | return None 69 | _id = str(uuid.uuid4()) 70 | data = { 71 | 'id': _id, 72 | 'image': utils.base64_encode(image), 73 | 'shape': image.shape 74 | } 75 | self._redis_db.rpush('image_queue', json.dumps(data)) 76 | 77 | result = None 78 | while result is None: 79 | time.sleep(0.01) 80 | result = self._redis_db.get(_id) 81 | 82 | result = self._decode_redis_data(result) 83 | result['id'] = _id 84 | if save: 85 | result['image'] = self._save_image(image, _id) 86 | 87 | if 'http' in image_path: 88 | result['url'] = image_path 89 | 90 | self._redis_db.delete(_id) 91 | return result 92 | -------------------------------------------------------------------------------- /imsearch/index.py: -------------------------------------------------------------------------------- 1 | """ 2 | index.py 3 | ==================================== 4 | The module contains the Index class 5 | """ 6 | 7 | import os 8 | import shutil 9 | import numpy as np 10 | import copy 11 | import tarfile 12 | import requests 13 | import pickle 14 | import uuid 15 | import tempfile 16 | 17 | from .nmslib import NMSLIBIndex, get_index_path 18 | from .extractor import FeatureExtractor 19 | from .repository import get_repository 20 | 21 | from .backend import run 22 | 23 | EPSILON = 10e-5 24 | 25 | 26 | class Index: 27 | """ 28 | The class to create the searchable index object. 29 | """ 30 | fe = None 31 | object_counter = 0 32 | 33 | def __init__(self, name): 34 | """ 35 | Create the index with its name. 36 | Parameters 37 | --------- 38 | name 39 | Unique indentifier for your searchable index object. 40 | """ 41 | Index.object_counter += 1 42 | if os.environ.get('DETECTOR_MODE') == 'local' and (Index.fe is None or Index.fe.poll() is not None): 43 | Index.fe = run() 44 | self.match_ratio = 0.3 45 | self.name = name 46 | self._nmslib_index = NMSLIBIndex(self.name) 47 | self._feature_extractor = FeatureExtractor(self.name) 48 | self._repository_db = get_repository(self.name, 'mongo') 49 | 50 | def __enter__(self): 51 | return self 52 | 53 | def __exit__(self, type, value, traceback): 54 | self._nmslib_index.createIndex() 55 | 56 | def __del__(self): 57 | Index.object_counter -= 1 58 | if Index.object_counter == 0 and Index.fe is not None: 59 | Index.fe.terminate() 60 | 61 | def _get_object_wise_similar(self, features, k): 62 | matches = {} 63 | for data in features['primary']: 64 | knn = self._nmslib_index.knnQuery( 65 | data['features'], 'primary', k=k) 66 | for x in knn: 67 | img_data = self._repository_db.find({ 68 | 'primary': { 69 | 'index': x[0], 70 | 'label': data['label'], 71 | 'name': data['name'] 72 | } 73 | }) 74 | if img_data is not None: 75 | _id = img_data['_id'] 76 | 77 | if _id in matches: 78 | matches[_id]['s_dist'] = x[1] + matches[_id]['s_dist'] 79 | else: 80 | matches[_id] = { 81 | 'data': copy.deepcopy(img_data), 82 | 's_dist': x[1] 83 | } 84 | 85 | knn = self._nmslib_index.knnQuery( 86 | features['secondary'], 'secondary', k=k) 87 | for x in knn: 88 | img_data = self._repository_db.find({'secondary_index': x[0]}) 89 | if img_data is not None: 90 | _id = img_data['_id'] 91 | if _id in matches: 92 | matches[_id]['p_dist'] = x[1] 93 | else: 94 | matches[_id] = { 95 | 'data': copy.deepcopy(img_data), 96 | 'p_dist': x[1] 97 | } 98 | 99 | matches = list(matches.values()) 100 | total_objects = float(len(features['primary'])) 101 | 102 | def update_scores(data): 103 | score = (1 - self.match_ratio)*data.get('p_dist', 0) / \ 104 | (total_objects + EPSILON) + \ 105 | self.match_ratio*data.get('s_dist', 0) 106 | return (data['data'], score) 107 | 108 | matches = list(map(update_scores, matches)) 109 | matches.sort(key=lambda x: x[1]) 110 | return matches 111 | 112 | def _get_metadata(self, features): 113 | metadata = {'primary': []} 114 | for obj in features['primary']: 115 | metadata['primary'].append({ 116 | 'name': obj['name'], 117 | 'box': obj['box'] 118 | }) 119 | 120 | return metadata 121 | 122 | @classmethod 123 | def loadFromFile(cls, path='imsearch_index.tar.gz', name=None): 124 | dir_path = os.path.join('/tmp', str(uuid.uuid4())) 125 | os.makedirs(dir_path, exist_ok=True) 126 | with tarfile.open(path, 'r:gz') as tf: 127 | tf.extractall(dir_path) 128 | 129 | with open(os.path.join(dir_path, 'info.pkl'), 'rb') as fp: 130 | info = pickle.load(fp) 131 | 132 | if name is not None: 133 | info['name'] = name 134 | 135 | with open(os.path.join(dir_path, 'data.pkl'), 'rb') as fp: 136 | data = pickle.load(fp) 137 | 138 | index_path = get_index_path(info['name']) 139 | os.makedirs(os.path.dirname(index_path), exist_ok=True) 140 | shutil.copyfile(os.path.join( 141 | dir_path, 'index', 'index.h5'), index_path) 142 | 143 | images_path = FeatureExtractor._get_image_path(info['name']) 144 | os.makedirs(images_path, exist_ok=True) 145 | updated_data = [] 146 | for entity in data: 147 | image = entity['image'] 148 | if image != '' and 'http' not in image: 149 | image_name = os.path.basename(image) 150 | dst = os.path.join(images_path, image_name) 151 | src = os.path.join(dir_path, 'images', image_name) 152 | shutil.copyfile(src, dst) 153 | entity['image'] = dst 154 | updated_data.append(copy.deepcopy(entity)) 155 | 156 | db = get_repository(info['name'], 'mongo') 157 | db.clean() 158 | db.insert(updated_data) 159 | 160 | return cls(info['name']) 161 | 162 | def saveIndex(self, path='imsearch_index.tar.gz'): 163 | data = self._repository_db.dump() 164 | info = { 165 | 'name': self.name, 166 | 'count': len(data) 167 | } 168 | with tarfile.open(path, 'w:gz') as tf: 169 | tf.add(get_index_path(self.name), 170 | arcname='index/index.h5', recursive=False) 171 | 172 | pkl_path = os.path.join('/tmp', "{}.pkl".format(str(uuid.uuid4()))) 173 | with open(pkl_path, 'wb') as temp: 174 | pickle.dump(data, temp) 175 | tf.add(pkl_path, arcname='data.pkl', recursive=False) 176 | 177 | with open(pkl_path, 'wb') as temp: 178 | pickle.dump(info, temp) 179 | tf.add(pkl_path, arcname='info.pkl', recursive=False) 180 | os.remove(pkl_path) 181 | 182 | tf.add(self._feature_extractor._get_image_path( 183 | self.name), arcname='images/') 184 | 185 | def cleanIndex(self): 186 | """ 187 | Cleans the index. It will delete the images already added to the index. It will also remove the database entry for the same index. 188 | """ 189 | self._feature_extractor.clean() 190 | self._nmslib_index.clean() 191 | self._repository_db.clean() 192 | 193 | def addImage(self, image_path, save=True): 194 | """ 195 | Add a single image to the index. 196 | Parameters 197 | --------- 198 | image_path 199 | The local path or url to the image to add to the index. 200 | """ 201 | features = self._feature_extractor.extract(image_path, save=save) 202 | if features is None: 203 | return False 204 | reposiory_data = self._nmslib_index.addDataPoint(features) 205 | self._repository_db.insert(reposiory_data) 206 | return True 207 | 208 | def addImageBatch(self, image_list, save=True): 209 | """ 210 | Add a multiple image to the index. 211 | Parameters 212 | --------- 213 | image_list 214 | The list of the image paths or urls to add to the index. 215 | """ 216 | response = [] 217 | for image_path in image_list: 218 | response.append(self.addImage(image_path, save=save)) 219 | return response 220 | 221 | def createIndex(self): 222 | """ 223 | Creates the index. Set create time paramenters and query-time parameters for nmslib index. 224 | """ 225 | self._nmslib_index.createIndex() 226 | 227 | def knnQuery(self, image_path, k=10, policy='global'): 228 | """ 229 | Query the index to search for k-nearest images in the database. 230 | Parameters 231 | --------- 232 | image_path 233 | The path to the query image. 234 | k=10 235 | Number of results. 236 | policy='global' 237 | choose policy from 'object' or 'global'. Search results will change accordingly. 238 | object: Object level matching. The engine will look for similarity at object level for every object detected in the image. 239 | global: Overall similarity using single feature space on the whole image. 240 | """ 241 | 242 | features = self._feature_extractor.extract(image_path, save=False) 243 | if features is None: 244 | return [] 245 | matches = [] 246 | if policy == 'object': 247 | matches = self._get_object_wise_similar(features, k) 248 | 249 | if not matches: 250 | knn = self._nmslib_index.knnQuery( 251 | features['secondary'], 'secondary', k=k) 252 | matches = [(self._repository_db.find( 253 | {'secondary_index': x[0]}, many=False), 1.0/(x[1] + EPSILON)) for x in knn] 254 | 255 | return matches[:k], self._get_metadata(features) 256 | -------------------------------------------------------------------------------- /imsearch/nmslib.py: -------------------------------------------------------------------------------- 1 | import os 2 | import shutil 3 | import pandas as pd 4 | import numpy as np 5 | import nmslib 6 | 7 | instance_map = {} 8 | 9 | 10 | def NMSLIBIndex(index_name): 11 | if index_name not in instance_map: 12 | instance_map[index_name] = _nmslibIndex(index_name) 13 | return instance_map[index_name] 14 | 15 | 16 | def get_index_path(index_name): 17 | home_dir = os.environ.get('HOME') 18 | return os.path.join(home_dir, '.imsearch', 'indices', index_name, 'index.h5') 19 | 20 | 21 | class _nmslibIndex: 22 | def __init__(self, name): 23 | self.index_name = name 24 | self._load_index() 25 | 26 | def _add_data(self, index, df): 27 | data = df.values.astype(np.float32) 28 | c = index.addDataPointBatch(data, range(data.shape[0])) 29 | return index, c 30 | 31 | def _load_index(self): 32 | index_file = get_index_path(self.index_name) 33 | 34 | self.primary = nmslib.init( 35 | method='hnsw', space='l2', data_type=nmslib.DataType.DENSE_VECTOR) 36 | self.secondary = nmslib.init( 37 | method='hnsw', space='l2', data_type=nmslib.DataType.DENSE_VECTOR) 38 | self.bitmap = nmslib.init( 39 | method='hnsw', space='l2', data_type=nmslib.DataType.DENSE_VECTOR) 40 | 41 | if os.path.exists(index_file): 42 | self.primary_df = pd.read_hdf(index_file, 'primary') 43 | self.primary, self.primary_c = self._add_data( 44 | self.primary, self.primary_df) 45 | 46 | self.secondary_df = pd.read_hdf(index_file, 'secondary') 47 | self.secondary, self.secondary_c = self._add_data( 48 | self.secondary, self.secondary_df) 49 | 50 | self.bitmap_df = pd.read_hdf(index_file, 'bitmap') 51 | self.bitmap, self.bitmap_c = self._add_data( 52 | self.bitmap, self.bitmap_df) 53 | else: 54 | self.primary_df = None 55 | self.secondary_df = None 56 | self.bitmap_df = None 57 | 58 | self.primary_c, self.secondary_c, self.bitmap_c = 0, 0, 0 59 | 60 | def _add_data_point(self, index, count, data): 61 | _t = index.addDataPoint(count, data) 62 | return index, _t + 1 63 | 64 | def clean(self): 65 | index_file = get_index_path(self.index_name) 66 | if os.path.exists(os.path.dirname(index_file)): 67 | shutil.rmtree(os.path.dirname(index_file)) 68 | self._load_index() 69 | 70 | def getDataByIds(self, id_list, _type='primary'): 71 | return [self.getDataById(x, _type) for x in id_list] 72 | 73 | def getDataById(self, _id, _type='primary'): 74 | return getattr(self, _type + '_df').iloc[_id].values 75 | 76 | def addDataPoint(self, data): 77 | primary = [] 78 | for d in data['primary']: 79 | primary.append({ 80 | 'index': self.primary_c, 81 | 'label': d['label'], 82 | 'name': d['name'], 83 | 'box': d['box'] 84 | }) 85 | v = d['features'] 86 | if self.primary_df is None: 87 | self.primary_df = pd.DataFrame(columns=range(v.shape[0])) 88 | self.primary_df.loc[self.primary_c] = v 89 | self.primary, self.primary_c = self._add_data_point( 90 | self.primary, self.primary_c, v) 91 | 92 | db_data = { 93 | '_id': data['id'], 94 | 'image': data['image'], 95 | 'url': data.get('url', ''), 96 | 'primary': primary, 97 | 'secondary_index': self.secondary_c, 98 | 'bitmap_index': self.bitmap_c 99 | } 100 | 101 | v = data['secondary'] 102 | if self.secondary_df is None: 103 | self.secondary_df = pd.DataFrame(columns=range(v.shape[0])) 104 | self.secondary_df.loc[self.secondary_c] = v 105 | self.secondary, self.secondary_c = self._add_data_point( 106 | self.secondary, self.secondary_c, v) 107 | 108 | v = np.array(data['object_bitmap'], dtype=np.float32) 109 | if self.bitmap_df is None: 110 | self.bitmap_df = pd.DataFrame(columns=range(v.shape[0])) 111 | self.bitmap_df.loc[self.bitmap_c] = v 112 | self.bitmap, self.bitmap_c = self._add_data_point( 113 | self.bitmap, self.bitmap_c, v) 114 | 115 | return db_data 116 | 117 | def addDataPointBatch(self, data_list): 118 | return [self.addDataPoint(data) for data in data_list] 119 | 120 | def createIndex(self): 121 | M = 15 122 | efC = 100 123 | num_threads = 4 124 | index_time_params = { 125 | 'M': M, 'indexThreadQty': num_threads, 'efConstruction': efC} 126 | self.primary.createIndex(index_time_params) 127 | self.secondary.createIndex(index_time_params) 128 | self.bitmap.createIndex(index_time_params) 129 | 130 | os.makedirs(os.path.dirname( 131 | get_index_path(self.index_name)), exist_ok=True) 132 | 133 | if self.primary_df is not None: 134 | self.primary_df.to_hdf(get_index_path(self.index_name), 'primary') 135 | 136 | if self.secondary_df is not None: 137 | self.secondary_df.to_hdf( 138 | get_index_path(self.index_name), 'secondary') 139 | 140 | if self.bitmap_df is not None: 141 | self.bitmap_df.to_hdf(get_index_path(self.index_name), 'bitmap') 142 | 143 | efs = 100 144 | query_time_params = {'efSearch': efs} 145 | self.primary.setQueryTimeParams(query_time_params) 146 | self.secondary.setQueryTimeParams(query_time_params) 147 | self.bitmap.setQueryTimeParams(query_time_params) 148 | 149 | def knnQuery(self, data, _type='bitmap', k=10): 150 | response = getattr(self, _type).knnQuery(data, k=k) 151 | return [(int(i), float(d)) for i, d in zip(response[0], response[1])] 152 | -------------------------------------------------------------------------------- /imsearch/repository/__init__.py: -------------------------------------------------------------------------------- 1 | from ..exception import InvalidAttributeError 2 | 3 | from .mongo import MongoRepository 4 | 5 | 6 | def get_repository(index_name, repo_type): 7 | repo = MongoRepository(index_name) 8 | return RepositoryWrapper(repo) 9 | 10 | 11 | class RepositoryWrapper: 12 | def __init__(self, repo): 13 | self.db = repo 14 | 15 | def clean(self): 16 | self.db.clean() 17 | 18 | def insert(self, data): 19 | if isinstance(data, dict): 20 | return self.db.insert_one(data) 21 | 22 | if isinstance(data, list): 23 | return self.db.insert_many(data) 24 | 25 | raise InvalidAttributeError( 26 | data, 'data of type dict or list expected, got {}'.format(type(data))) 27 | 28 | def find(self, query, many=False): 29 | if many: 30 | return self.db.find(query) 31 | else: 32 | return self.db.find_one(query) 33 | 34 | def dump(self): 35 | return self.db.find({}) 36 | -------------------------------------------------------------------------------- /imsearch/repository/mongo.py: -------------------------------------------------------------------------------- 1 | import os 2 | from pymongo import MongoClient 3 | 4 | 5 | class MongoRepository: 6 | def __init__(self, index_name): 7 | url = os.environ.get('MONGO_URI') 8 | self.db = MongoClient(url).imsearch[index_name] 9 | 10 | def clean(self): 11 | self.db.drop() 12 | 13 | def insert_one(self, data): 14 | return self.db.insert_one(data) 15 | 16 | def insert_many(self, data): 17 | return self.db.insert_many(data) 18 | 19 | def find_one(self, query): 20 | response = self.db.find_one(query) 21 | return response 22 | 23 | def find(self, query): 24 | return list(self.db.find(query)) 25 | -------------------------------------------------------------------------------- /imsearch/storage/__init__.py: -------------------------------------------------------------------------------- 1 | import requests 2 | 3 | from .gcloud import GcloudStorage 4 | from .aws import S3Storage 5 | 6 | def get_storage_object(service_name='gcp'): 7 | if service_name == 'gcp': 8 | return GcloudStorage() 9 | 10 | if service_name == 's3': 11 | return S3Storage() 12 | 13 | return None -------------------------------------------------------------------------------- /imsearch/storage/abstract.py: -------------------------------------------------------------------------------- 1 | class AbstractStorage(): 2 | def upload(): 3 | raise NotImplementedError("Subclass should implement this method") -------------------------------------------------------------------------------- /imsearch/storage/aws.py: -------------------------------------------------------------------------------- 1 | import os 2 | import boto3 3 | from .abstract import AbstractStorage 4 | 5 | 6 | class S3Storage(AbstractStorage): 7 | def __init__(self): 8 | AWS_ACCESS_KEY_ID = os.environ.get('AWS_ACCESS_KEY_ID', False) 9 | AWS_SECRET_ACCESS_KEY = os.environ.get('AWS_SECRET_ACCESS_KEY', False) 10 | if not (AWS_ACCESS_KEY_ID and AWS_SECRET_ACCESS_KEY): 11 | raise Exception( 12 | 'AWS s3 credential config not provided: AWS_ACCESS_KEY_ID, AWS_SECRET_ACCESS_KEY') 13 | 14 | self.bucket_name = os.environ.get('BUCKET_NAME', 'imsearch') 15 | self.region = os.environ.get('S3_REGION', 'us-east-2') 16 | self.client = boto3.client('s3', region_name=self.region, 17 | aws_access_key_id=AWS_ACCESS_KEY_ID, aws_secret_access_key=AWS_SECRET_ACCESS_KEY) 18 | 19 | try: 20 | self.client.create_bucket(Bucket=self.bucket_name, CreateBucketConfiguration={ 21 | 'LocationConstraint': self.region}) 22 | except self.client.exceptions.BucketAlreadyExists: 23 | print('{}: Bucket already exists'.format(self.bucket_name)) 24 | except self.client.exceptions.BucketAlreadyOwnedByYou: 25 | print('{}: Bucket already owned by you'.format(self.bucket_name)) 26 | 27 | def upload(self, image_path, key): 28 | content_type = 'image/{}'.format(key.split('.')[-1]) 29 | self.client.upload_file(image_path, self.bucket_name, key, ExtraArgs={ 30 | 'ACL': 'public-read', 'ContentType': content_type}) 31 | 32 | return 'https://s3.{}.amazonaws.com/{}/{}'.format(self.region, self.bucket_name, key) 33 | -------------------------------------------------------------------------------- /imsearch/storage/gcloud.py: -------------------------------------------------------------------------------- 1 | import os 2 | from google.cloud import storage 3 | from .abstract import AbstractStorage 4 | 5 | 6 | class GcloudStorage(AbstractStorage): 7 | def __init__(self): 8 | if not os.environ.get('GOOGLE_APPLICATION_CREDENTIALS', False): 9 | raise Exception( 10 | 'Google cloud credential config not provided: GOOGLE_APPLICATION_CREDENTIALS') 11 | 12 | self.bucket_name = os.environ.get('BUCKET_NAME', 'imsearch') 13 | self.client = storage.Client() 14 | 15 | try: 16 | bucket = self.client.lookup_bucket(self.bucket_name) 17 | if not bucket: 18 | self.bucket = self.client.create_bucket(self.bucket_name) 19 | else: 20 | self.bucket = bucket 21 | except Exception as e: 22 | raise e 23 | 24 | def upload(self, image_path, key): 25 | blob = self.bucket.blob(key) 26 | blob.upload_from_filename(image_path) 27 | blob.make_public() 28 | return blob.public_url 29 | -------------------------------------------------------------------------------- /imsearch/utils/__init__.py: -------------------------------------------------------------------------------- 1 | from .image import * 2 | -------------------------------------------------------------------------------- /imsearch/utils/image.py: -------------------------------------------------------------------------------- 1 | import sys 2 | import base64 3 | import requests 4 | import cv2 5 | from PIL import Image 6 | from io import BytesIO 7 | 8 | import numpy as np 9 | 10 | 11 | def _get_data_from_path(image_path): 12 | if image_path.startswith('data:image/'): 13 | image_path = image_path.split(',')[1] 14 | data = BytesIO(base64.urlsafe_b64decode(image_path)) 15 | elif image_path.startswith('http'): 16 | r = requests.get(image_path) 17 | if r.status_code != 200: 18 | return None 19 | data = BytesIO(r.content) 20 | else: 21 | data = image_path 22 | 23 | return data 24 | 25 | 26 | def check_load_image(image_path): 27 | if isinstance(image_path, str): 28 | data = _get_data_from_path(image_path) 29 | try: 30 | img = np.asarray(Image.open(data)) 31 | except: 32 | return None 33 | else: 34 | try: 35 | img = np.asarray(Image.open(image_path)) 36 | except: 37 | return None 38 | 39 | if len(img.shape) != 3: 40 | return None 41 | 42 | if img.shape[2] > 3: 43 | img = img[:, :, :3] 44 | 45 | return img 46 | 47 | 48 | def save_image(image, dest): 49 | image = Image.fromarray(image) 50 | image.save(dest) 51 | 52 | 53 | def base64_encode(a): 54 | a = a.copy(order='C') 55 | return base64.b64encode(a).decode("utf-8") 56 | 57 | 58 | def base64_decode(a, shape=None, dtype=np.uint8): 59 | if sys.version_info.major == 3: 60 | a = bytes(a, encoding="utf-8") 61 | 62 | a = np.frombuffer(base64.decodestring(a), dtype=dtype) 63 | if shape is not None: 64 | a = a.reshape(shape) 65 | return a 66 | -------------------------------------------------------------------------------- /install.sh: -------------------------------------------------------------------------------- 1 | echo "Installing nmslib" 2 | sudo apt-get install python3-dev 3 | pip install --no-binary :all: nmslib 4 | 5 | echo "Installing imsearch" 6 | pip install . -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- 1 | Cython 2 | requests 3 | numpy 4 | torch>=1.4.0 5 | matplotlib 6 | wget 7 | opencv_python 8 | boto3 9 | redis 10 | pandas 11 | torchvision 12 | pymongo 13 | google-cloud-storage 14 | -------------------------------------------------------------------------------- /setup.py: -------------------------------------------------------------------------------- 1 | import setuptools 2 | 3 | with open("README.md", "r") as fh: 4 | long_description = fh.read() 5 | 6 | setuptools.setup( 7 | name='imsearch', 8 | version='0.1.1', 9 | description='A generic framework to build your own reverse image search engine', 10 | long_description=long_description, 11 | long_description_content_type="text/markdown", 12 | url='https://github.com/rikenmehta03/imsearch', 13 | author='Riken Mehta', 14 | author_email='riken.mehta03@gmail.com', 15 | packages=setuptools.find_packages(), 16 | install_requires=['Cython', 'torch', 'torchvision', 'pandas', 'tables', 'redis', 'pymongo', 'nmslib', 'wget', 'opencv-python', 'requests', 'google-cloud-storage', 'boto3'], 17 | classifiers=[ 18 | "Programming Language :: Python :: 3", 19 | "License :: OSI Approved :: Apache Software License", 20 | "Operating System :: POSIX :: Linux", 21 | ] 22 | ) 23 | -------------------------------------------------------------------------------- /tests/test_index.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | from imsearch import Index 3 | import imsearch 4 | 5 | import unittest 6 | from unittest.mock import patch, Mock 7 | 8 | 9 | class TestIndex(unittest.TestCase): 10 | @patch('imsearch.index.NMSLIBIndex') 11 | @patch('imsearch.index.FeatureExtractor') 12 | @patch('imsearch.index.get_repository') 13 | @patch('imsearch.index.os.environ') 14 | def setUp(self, mock_os_environ, mock_get_repository, mock_extractor, mock_nmslib): 15 | self.mock_nmslib_instance = mock_nmslib() 16 | self.mock_extractor_instance = mock_extractor() 17 | self.mock_get_repository = mock_get_repository 18 | self.mock_repository_instance = Mock() 19 | self.mock_get_repository.return_value = self.mock_repository_instance 20 | mock_os_environ.get.return_value = 'remote' 21 | 22 | self.test_index = Index('test') 23 | 24 | def test_constructor(self): 25 | self.assertIsInstance(self.test_index, Index) 26 | self.assertEqual(self.test_index.name, 'test') 27 | self.assertEqual(self.test_index._nmslib_index, 28 | self.mock_nmslib_instance) 29 | self.assertEqual(self.test_index._feature_extractor, 30 | self.mock_extractor_instance) 31 | self.mock_get_repository.assert_called_with('test', 'mongo') 32 | 33 | def test_add_image_valid(self): 34 | IMG_PATH = '../images/000000000139.jpg' 35 | mock_data = { 36 | 'id': '12345', 37 | 'image': IMG_PATH, 38 | 'secondary': np.zeros(1024) 39 | } 40 | 41 | self.mock_extractor_instance.extract.return_value = mock_data 42 | self.mock_nmslib_instance.addDataPoint.return_value = mock_data 43 | 44 | self.assertEqual(self.test_index.addImage(IMG_PATH), True) 45 | self.mock_extractor_instance.extract.assert_called_with(IMG_PATH, save=True) 46 | self.mock_nmslib_instance.addDataPoint.assert_called_with(mock_data) 47 | self.mock_repository_instance.insert.assert_called_with(mock_data) 48 | 49 | def test_add_image_invalid(self): 50 | IMG_PATH = '../images/000000000139.jpg' 51 | 52 | self.mock_extractor_instance.extract.return_value = None 53 | 54 | self.assertEqual(self.test_index.addImage(IMG_PATH), False) 55 | self.mock_extractor_instance.extract.assert_called_with(IMG_PATH, save=True) 56 | self.mock_nmslib_instance.addDataPoint.assert_not_called() 57 | self.mock_repository_instance.insert.assert_not_called() 58 | 59 | def test_add_image_batch(self): 60 | list_length = 20 61 | IMG_PATH = '../images/000000000139.jpg' 62 | IMAGE_LIST = [IMG_PATH] * list_length 63 | mock_data = { 64 | 'id': '12345', 65 | 'image': IMG_PATH, 66 | 'secondary': np.zeros(1024) 67 | } 68 | 69 | self.mock_extractor_instance.extract.return_value = mock_data 70 | self.mock_nmslib_instance.addDataPoint.return_value = mock_data 71 | 72 | self.assertEqual(self.test_index.addImageBatch( 73 | IMAGE_LIST), [True]*list_length) 74 | self.assertEqual( 75 | self.mock_extractor_instance.extract.call_count, list_length) 76 | self.assertEqual( 77 | self.mock_nmslib_instance.addDataPoint.call_count, list_length) 78 | self.assertEqual( 79 | self.mock_repository_instance.insert.call_count, list_length) 80 | 81 | def test_create_index(self): 82 | self.test_index.createIndex() 83 | self.mock_nmslib_instance.createIndex.assert_called() 84 | 85 | def test_knn_query(self): 86 | pass 87 | 88 | 89 | if __name__ == "__main__": 90 | unittest.main() 91 | -------------------------------------------------------------------------------- /tests/test_init.py: -------------------------------------------------------------------------------- 1 | import imsearch 2 | 3 | import unittest 4 | from unittest.mock import patch 5 | 6 | 7 | class TestInit(unittest.TestCase): 8 | 9 | @patch('imsearch.config_init') 10 | @patch('imsearch.Index') 11 | def test_init(self, mock_index, config_init_mock): 12 | instance = mock_index() 13 | instance.name = 'test' 14 | 15 | index = imsearch.init('test', MONGO_URI='mongodb://localhost:27017/') 16 | self.assertEqual(index, instance) 17 | self.assertEqual(index.name, 'test') 18 | config_init_mock.assert_called_with( 19 | {'MONGO_URI': 'mongodb://localhost:27017/'}) 20 | 21 | @patch('imsearch.run') 22 | def test_detector(self, mock_backend): 23 | imsearch.run_detector('redis://dummy:Password@111.111.11.111:6379/0') 24 | mock_backend.assert_called_with( 25 | 'redis://dummy:Password@111.111.11.111:6379/0') 26 | 27 | 28 | if __name__ == "__main__": 29 | unittest.main() 30 | --------------------------------------------------------------------------------