import sys
import torch
import clip
import psycopg2
from PIL import Image

device = "cpu"

model, preprocess = clip.load("ViT-B/32", device=device)

image = preprocess(Image.open(sys.argv[1])).unsqueeze(0).to(device)

with torch.no_grad():
    vec = model.encode_image(image)

vec = vec / vec.norm(dim=-1, keepdim=True)

vector = vec.cpu().numpy()[0].tolist()
vector = "[" + ",".join(map(str, vector)) + "]"

conn = psycopg2.connect(
    host="192.168.1.42",
    database="postgres",
    user="postgres",
    password="yiparts007"
)

cur = conn.cursor()

cur.execute("""
SELECT path
FROM images
ORDER BY embedding <-> %s::vector
LIMIT 10
""", (vector,))

rows = cur.fetchall()

for r in rows:
    print(r[0])