Spaces:

Supabase
/

image-search

Paused

App Files Files Community

Andrew Smith commited on Jul 4, 2023

Commit

48e00e6

1 Parent(s): ad7257f

Initial commit

Browse files

Files changed (13) hide show

.gitignore +3 -0
Dockerfile +18 -0
image_search/main.py +101 -0
images/four.jpg +0 -0
images/one.jpg +0 -0
images/three.jpg +0 -0
images/two.jpg +0 -0
poetry.lock +0 -0
pyproject.toml +23 -0
static/index.html +21 -0
static/script.js +47 -0
static/style.css +9 -0
tests/__init__.py +0 -0

.gitignore ADDED Viewed

	@@ -0,0 +1,3 @@

+__pycache__
+.vscode
+supabase/

Dockerfile ADDED Viewed

	@@ -0,0 +1,18 @@

+FROM python:3.10
+WORKDIR /code
+COPY . /code
+# Get secret DB_URL and clone it as repo at buildtime
+RUN --mount=type=secret,id=DB_URL,mode=0444,required=true \
+  git clone $(cat /run/secrets/DB_URL)
+RUN pip install --no-cache-dir poetry
+RUN poetry config virtualenvs.create false \
+ && poetry install --no-dev
+RUN poetry run search
+CMD ["poetry", "run", "start"]

image_search/main.py ADDED Viewed

	@@ -0,0 +1,101 @@

+import os
+from typing import Union
+from PIL import Image
+from fastapi import FastAPI
+from fastapi.staticfiles import StaticFiles
+from fastapi.responses import FileResponse
+from sentence_transformers import SentenceTransformer
+import uvicorn
+import vecs
+DB_CONNECTION = os.environ.get(
+    'DB_URL', "postgresql://postgres:postgres@localhost:54322/postgres")
+app = FastAPI()
+@app.get("/seed")
+def seed():
+    # create vector store client
+    vx = vecs.create_client(DB_CONNECTION)
+    iv = vx.get_collection(name="image_vectors")
+    if iv:
+        return {"message": "Collection already exists."}
+    # create a collection of vectors with 512 dimensions
+    images = vx.create_collection(name="image_vectors", dimension=512)
+    # Load CLIP model
+    model = SentenceTransformer('clip-ViT-B-32')
+    # Encode an image:
+    img_emb1 = model.encode(Image.open('./images/one.jpg'))
+    img_emb2 = model.encode(Image.open('./images/two.jpg'))
+    img_emb3 = model.encode(Image.open('./images/three.jpg'))
+    img_emb4 = model.encode(Image.open('./images/four.jpg'))
+    images.upsert(
+        vectors=[
+            (
+                "one.jpg",
+                img_emb1,
+                {"type": "jpg"}
+            ), (
+                "two.jpg",
+                img_emb2,
+                {"type": "jpg"}
+            ), (
+                "three.jpg",
+                img_emb3,
+                {"type": "jpg"}
+            ), (
+                "four.jpg",
+                img_emb4,
+                {"type": "jpg"}
+            )
+        ]
+    )
+    print("Inserted images")
+    # index the collection fro fast search performance
+    images.create_index()
+    return {"message": "Collection created and indexed."}
+@app.get("/search")
+def search(query: Union[str, None] = None):
+    # create vector store client
+    vx = vecs.create_client(DB_CONNECTION)
+    images = vx.get_collection(name="image_vectors")
+    # Load CLIP model
+    model = SentenceTransformer('clip-ViT-B-32')
+    # Encode text query
+    query_string = query
+    text_emb = model.encode(query_string)
+    # query the collection filtering metadata for "type" = "jpg"
+    results = images.query(
+        query_vector=text_emb,
+        limit=1,
+        filters={"type": {"$eq": "jpg"}},
+    )
+    result = results[0]
+    return {"result": result, "query": query}
+app.mount("/images", StaticFiles(directory="images"), name="images")
+app.mount("/", StaticFiles(directory="static", html=True), name="static")
+@app.get("/")
+def index() -> FileResponse:
+    return FileResponse(path="static/index.html", media_type="text/html")
+def start():
+    """Launched with `poetry run start` at root level"""
+    uvicorn.run("image_search.main:app",
+                host="0.0.0.0", port=7860, reload=True)

images/four.jpg ADDED Viewed

images/one.jpg ADDED Viewed

images/three.jpg ADDED Viewed

images/two.jpg ADDED Viewed

poetry.lock ADDED Viewed

The diff for this file is too large to render. See raw diff

pyproject.toml ADDED Viewed

	@@ -0,0 +1,23 @@

+[tool.poetry]
+name = "image-search"
+version = "0.1.0"
+description = ""
+authors = ["Andrew Smith <[email protected]>"]
+readme = "README.md"
+packages = [{include = "image_search"}]
+[tool.poetry.dependencies]
+python = "^3.10"
+vecs = "^0.2.6"
+sentence-transformers = "^2.2.2"
+fastapi = "^0.99.1"
+uvicorn = {extras = ["standard"], version = "^0.22.0"}
+[tool.poetry.scripts]
+start = "image_search.main:start"
+seed = "image_search.main:seed"
+search = "image_search.main:search"
+[build-system]
+requires = ["poetry-core"]
+build-backend = "poetry.core.masonry.api"

static/index.html ADDED Viewed

	@@ -0,0 +1,21 @@

+<!DOCTYPE html>
+<html lang="en">
+<head>
+    <meta charset="UTF-8">
+    <meta name="viewport" content="width=device-width, initial-scale=1.0">
+    <title>Image Search</title>
+    <link rel="stylesheet" href="style.css" />
+    <script type="module" src="script.js"></script>
+</head>
+<body>
+    <main>
+        <h1>Image Search</h1>
+        <form class="search-form">
+            <label>Name </label>
+            <input id="query" type="text" name="query">
+            <button class="submit-button">Submit</button>
+        </form>
+        <div class="result"></div>
+    </main>
+</body>
+</html>

static/script.js ADDED Viewed

	@@ -0,0 +1,47 @@

+const searchForm = document.querySelector(".search-form");
+const result = document.querySelector(".result");
+const submitButton = document.querySelector(".submit-button");
+const clearDiv = () => {
+  result.innerHTML = "";
+};
+const buttonLoading = () => {
+  submitButton.innerHTML = "loading...";
+  submitButton.setAttribute("disabled", true);
+};
+const buttonDefault = () => {
+  submitButton.innerHTML = "Submit";
+  submitButton.removeAttribute("disabled");
+};
+const renderImg = (imageName) => {
+  clearDiv();
+  const img = document.createElement("img");
+  img.src = `/images/${imageName}`;
+  result.appendChild(img);
+};
+const getResult = async (searchQuery) => {
+  const response = await fetch(`search?query=${searchQuery}`);
+  const json = await response.json();
+  return json.result;
+};
+searchForm.addEventListener("submit", async (event) => {
+  event.preventDefault();
+  const query = document.querySelector("#query");
+  buttonLoading();
+  try {
+    const res = await getResult(query.value);
+    renderImg(res);
+  } catch (err) {
+    console.error(err);
+    clearDiv();
+  } finally {
+    buttonDefault();
+  }
+});

static/style.css ADDED Viewed

	@@ -0,0 +1,9 @@

+main {
+  display: flex;
+  flex-direction: column;
+  justify-content: center;
+  align-items: center;
+}
+.search-form {
+    padding: 20px 0;
+}

tests/__init__.py ADDED Viewed

File without changes