Initial Commit

Author: ItzCrazyKns

.eslintrc.json

@@ -0,0 +1,3 @@
+{
+  "extends": "next/core-web-vitals"
+}

.example.env.local

@@ -0,0 +1,4 @@
+DATABASE_URL=
+GOOGLE_CLOUD_PROJECT=
+GOOGLE_CLOUD_LOCATION=
+GOOGLE_CLOUD_API_ENDPOINT=

.gitignore

@@ -0,0 +1,36 @@
+# See https://help.github.com/articles/ignoring-files/ for more about ignoring files.
+
+# dependencies
+/node_modules
+/.pnp
+.pnp.js
+.yarn/install-state.gz
+
+# testing
+/coverage
+
+# next.js
+/.next/
+/out/
+
+# production
+/build
+
+# misc
+.DS_Store
+*.pem
+
+# debug
+npm-debug.log*
+yarn-debug.log*
+yarn-error.log*
+
+# local env files
+.env*.local
+
+# vercel
+.vercel
+
+# typescript
+*.tsbuildinfo
+next-env.d.ts

README.md

@@ -0,0 +1,55 @@
+# Neon Image Search
+
+This is an image search app, built using NeonDB and Vertex AI. It uses Vertex AI to convert the data images into embeddings and then stores it into NeonDB. When a request is made the request image is converted into embeddings and then the similarity search is performed by NeonDB using PGVector.
+
+![preview](./assets//app_preview.png)
+
+## Installation
+
+Before we began with the installation make sure you're logged in, in the [Google Cloud CLI](https://cloud.google.com/docs/authentication/gcloud#gcloud-credentials).
+
+1. Clone the repo
+
+```bash
+git clone https://github.com/ItzCrazyKns/Neon-Image-Search.git
+```
+
+2. Install dependencies
+
+```bash
+yarn
+```
+
+Or if you are using NPM
+
+```bash
+npm i
+```
+
+3. Change the name of `.example.env.local` to `.env.local` and fill all the fields.
+
+4. Convert data images into embeddings. (The data images are stored in `/public/flower_images`)
+
+```bash
+yarn run generate-embeddings
+```
+
+Or if you are using NPM
+
+```bash
+npm run generate-embeddings
+```
+
+5. Then finally started the app
+
+```bash
+yarn run dev
+```
+
+Or if you are using NPM
+
+```bash
+npm run dev
+```
+
+**Note**: You can build the project then also use it by following NextJS's guidelines.

app/api/predict/route.ts

@@ -0,0 +1,60 @@
+import { NextRequest, NextResponse } from "next/server";
+import { neon, neonConfig } from '@neondatabase/serverless';
+import { PredictionServiceClient } from "@google-cloud/aiplatform";
+
+neonConfig.fetchConnectionCache = true;
+
+const sql = neon(process.env.DATABASE_URL!);
+
+const vertexAiClient = new PredictionServiceClient({
+    apiEndpoint: process.env.GOOGLE_CLOUD_API_ENDPOINT,
+})
+
+export const POST = async (req: NextRequest) => {
+    try {
+        const data = await req.json()
+        
+        if (!data.base64Image) return NextResponse.json({ error: 'No image provided' }, { status: 400 })
+
+        const response = await vertexAiClient.predict({
+            endpoint: `projects/${process.env.GOOGLE_CLOUD_PROJECT}/locations/${process.env.GOOGLE_CLOUD_LOCATION}/publishers/google/models/multimodalembedding@001`,
+            instances: [
+                {
+                    structValue: {
+                        fields: {
+                            image: {
+                                structValue: {
+                                    fields: {
+                                        bytesBase64Encoded: {
+                                            stringValue: data.base64Image,
+                                        },
+                                    },
+                                },
+                            },
+                        },
+                    },
+                },
+            ],
+        });
+
+        if (!response || response.length < 1 || !response[0]?.predictions || response[0]?.predictions.length < 1) return NextResponse.json({ error: 'No predictions found for image' }, { status: 404 })
+
+        const embeddings = response[0]?.predictions[0]?.structValue?.fields?.imageEmbedding?.listValue?.values?.map((v: any) => v.numberValue)
+
+        const result = await sql(`SELECT  id, type, name FROM items ORDER BY embedding::VECTOR <=> '[${embeddings}]' LIMIT 8;`)
+
+        const similarImages = result.map((item: any) => {
+            return {
+                id: item.id,
+                type: item.type,
+                name: item.name,
+                image: '/flower_images/' + item.type + '/' + item.name,
+            }
+        })
+
+        return NextResponse.json({ result: similarImages }, { status: 200 })
+    } catch (err) {
+        console.error('An error has occurred:', err)
+        return NextResponse.json({ error: 'An error has occurred' }, { status: 500 })
+    }
+}

app/favicon.ico

@@ -0,0 +1,3 @@
+@tailwind base;
+@tailwind components;
+@tailwind utilities;

app/globals.css

@@ -0,0 +1,19 @@
+import type { Metadata } from "next";
+import "./globals.css";
+
+export const metadata: Metadata = {
+  title: "Image Search Engine",
+  description: "Search for similar images using Neon DB and Vertex AI",
+};
+
+export default function RootLayout({
+  children,
+}: Readonly<{
+  children: React.ReactNode;
+}>) {
+  return (
+    <html lang="en">
+      <body>{children}</body>
+    </html>
+  );
+}

app/layout.tsx

@@ -0,0 +1,75 @@
+'use client';
+
+import Image from 'next/image';
+import React, { useCallback, useState } from 'react';
+import { useDropzone } from 'react-dropzone';
+
+const Home = () => {
+  const [filePreview, setFilePreview] = useState("");
+  const [base64Image, setBase64Image] = useState("");
+  const [similarImages, setSimilarImages] = useState([] as any[]);
+  const [loading, setLoading] = useState(false)
+
+  const onDrop = useCallback(async (acceptedFiles: File[]) => {
+    const file = acceptedFiles[0]
+    const previewUrl = URL.createObjectURL(file)
+    setFilePreview(previewUrl)
+    const arrayBuffer = await file.arrayBuffer();
+    setBase64Image(Buffer.from(arrayBuffer).toString('base64'))
+  }, []);
+
+  const { getRootProps, getInputProps, isDragActive } = useDropzone({
+    onDrop,
+    multiple: false,
+  });
+
+  return (
+    <main className='bg-gray-950 min-h-screen flex flex-col items-center justify-center'>
+      <div className='max-w-screen-xl m-8 xl:m-auto flex flex-col items-center justify-center'>
+        <h2 className='text-lg sm:text-xl md:text-2xl xl:text-4xl font-semibold text-white text-center mt-8'><span className='bg-gradient-to-br from-green-600 to-green-300 bg-clip-text text-transparent p-1'>Search</span> for similar images 🚀</h2>
+        <p className='text-gray-300 text-xs md:text-sm max-w-[300px] text-center mt-0 xl:mt-2 2xl:mt-3'>Upload an image of Lilly, Lotus, Orchid, Sunflower, or Tulip and find images similar to it.</p>
+
+        <div className='text-white border border-white/30 rounded-lg cursor-pointer mt-8' {...getRootProps()}>
+          <input {...getInputProps()} />
+          {
+            isDragActive ?
+              <p className='my-10 mx-5 text-center'>Drop the file here ...</p> :
+              filePreview ?
+                <div>
+                  <Image width={512} height={512} src={filePreview} alt="Preview" className='h-[200px] w-[200px] md:h-[300px] md:w-auto m-1 md:m-2 xl:m-4 rounded-lg' />
+                </div> :
+                <p className='my-10 mx-5 text-center text-xs lg:text-lg'>Drag &apos;n&apos; drop a file here, or click to select a file</p>
+          }
+        </div>
+        <button
+          onClick={async () => {
+            setLoading(true)
+            const res = await fetch('/api/predict', {
+              method: "POST",
+              body: JSON.stringify({
+                base64Image: base64Image
+              })
+            })
+
+            const data = await res.json()
+            setSimilarImages(data.result)
+            setLoading(false)
+          }}
+          disabled={loading}
+          className='bg-gradient-to-br disabled:opacity-70 from-green-600 to-green-300 border border-green-600 rounded-lg w-full md:w-2/4 xl:w-4/12 my-4 text-white font-semibold text-sm xl:text-lg py-2 hover:opacity-90 transition duration-200'
+        >Search
+        </button>
+        <div className='grid grid-cols-1 sm:grid-cols-2 md:grid-cols-3 gap-4 mb-8'>
+          {similarImages.map((image, index) => (
+            <div key={index} className='flex flex-col items-center border border-white/30 rounded-lg pb-2'>
+              <Image width={512} height={512} src={image.image} alt={image.name} className='h-[200px] w-[200px] md:h-[250px] md:w-auto m-1 md:m-2 xl:m-4 rounded-lg' />
+              <p className='text-white font-semibold text-lg xl:-mt-2'>{image.type}</p>
+            </div>
+          ))}
+        </div>
+      </div>
+    </main>
+  );
+}
+
+export default Home;

app/page.tsx

@@ -0,0 +1,120 @@
+import { Pool } from 'pg';
+import * as fs from "fs/promises"
+import * as path from 'path';
+import dotenv from 'dotenv';
+import { PredictionServiceClient } from '@google-cloud/aiplatform';
+
+dotenv.config({
+    path: path.join(__dirname, '.env.local')
+});
+
+const vertexAiClient = new PredictionServiceClient({
+    apiEndpoint: process.env.GOOGLE_CLOUD_API_ENDPOINT,
+})
+
+const pool = new Pool({
+    connectionString: process.env.DATABASE_URL
+})
+
+pool.connect()
+
+const baseDir = "./public/flower_images";
+
+const listFolders = async (baseDir: string): Promise<string[]> => {
+    const items = await fs.readdir(baseDir, { withFileTypes: true });
+    return items.filter(item => item.isDirectory()).map(item => path.join(baseDir, item.name));
+};
+
+const listImages = async (folderPath: string): Promise<{ name: string; path: string; folderPath: string; }[]> => {
+    const items = await fs.readdir(folderPath, { withFileTypes: true })
+
+    return items.filter(item => item.isFile() && (item.name.includes('.jpg') || item.name.includes('.jpeg') || item.name.includes('.png'))).slice(0, 50).map(item => {
+        return {
+            path: path.join(folderPath, item.name),
+            folderPath: folderPath.split('\\').pop()!,
+            name: item.name
+        }
+    })
+}
+
+const imageToBase64 = async (filePath: string): Promise<string> => {
+    const fileBuffer = await fs.readFile(filePath)
+    return fileBuffer.toString('base64')
+}
+
+const convertImagesToEmbeddings = async (): Promise<any[]> => {
+    const embeddings: any[] = []
+    const folders = await listFolders(baseDir)
+
+    for await (const folder of folders) {
+        const images = await listImages(folder)
+        for await (const image of images) {
+            const base64Image = await imageToBase64(image.path)
+
+            const response = await vertexAiClient.predict({
+                endpoint: `projects/${process.env.GOOGLE_CLOUD_PROJECT}/locations/${process.env.GOOGLE_CLOUD_LOCATION}/publishers/google/models/multimodalembedding@001`,
+                instances: [
+                    {
+                        structValue: {
+                            fields: {
+                                image: {
+                                    structValue: {
+                                        fields: {
+                                            bytesBase64Encoded: {
+                                                stringValue: base64Image,
+                                            },
+                                        },
+                                    },
+                                },
+                            },
+                        },
+                    },
+                ],
+            });
+
+            if (!response || response.length < 1 || !response[0]?.predictions || response[0]?.predictions.length < 1) {
+                console.error('No predictions found for image', image.path);
+                continue;
+            }
+
+            embeddings.push({
+                embeddings: response[0]?.predictions[0]?.structValue?.fields?.imageEmbedding?.listValue?.values?.map((v) => v.numberValue),
+                type: image.folderPath.split('\\').pop()!,
+                name: image.name
+            })
+        }
+    }
+
+    return embeddings
+};
+
+const generateEmbeddings = async () => {
+    const client = await pool.connect()
+    try {
+        await client.query(`
+        CREATE EXTENSION IF NOT EXISTS vector;
+        CREATE TABLE IF NOT EXISTS items (
+            id SERIAL PRIMARY KEY,
+            type VARCHAR(255) NOT NULL,
+            name VARCHAR(255) NOT NULL,
+            embedding VECTOR(1408) NOT NULL
+        )
+        `)
+        console.log('Successfully created tables')
+        console.log('Generating embeddings...')
+
+        const embeddings = await convertImagesToEmbeddings()
+        console.log(embeddings)
+        console.log('Inserting embeddings into database...')
+        await client.query(`
+        INSERT INTO items (type, name, embedding) VALUES ${embeddings.map((e) => `('${e.type}', '${e.name}', '[${e.embeddings}]')`).join(',')}
+        `)
+        console.log('Successfully inserted embeddings into database')
+    } catch (err) {
+        console.error('An error has occurred:', err)
+    } finally {
+        client.release()
+    }
+}
+
+generateEmbeddings()

assets/app_preview.png

@@ -0,0 +1,7 @@
+import { PredictionServiceClient } from "@google-cloud/aiplatform";
+
+const vertexAiClient = new PredictionServiceClient({
+    apiEndpoint: process.env.GOOGLE_CLOUD_API_ENDPOINT,
+})
+
+export default vertexAiClient