import os
import torch
import clip
import psycopg2
from PIL import Image

device = "cuda" if torch.cuda.is_available() else "cpu"

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

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

cur = conn.cursor()

img_dir = "../images"

for file in os.listdir(img_dir):

    path = os.path.join(img_dir, file)

    image = preprocess(Image.open(path)).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()

    cur.execute(
        "INSERT INTO images(path, embedding) VALUES (%s,%s)",
        (file, vector)
    )

    print("insert:", file)

conn.commit()