Query by Text

# code for testing
host = "<your opensearch domain>"
port = 443
user = "<your master user>"
password = "<your master password>"

client = OpenSearch(
    hosts=[{"host": host, "port": port}],
    http_auth=(user, password),
    use_ssl=True,
    verify_certs=True,
    connection_class=RequestsHttpConnection,
    timeout = 60,
    # max_retries = 5,
    # retry_on_timeout = True
)

s3 = boto3.client("s3")
BUCKET_NAME = "ai-challenge-2024"
FOLDER_NAME = "keyframes"

model = SentenceTransformer("clip-ViT-B-32")

def query(client , model ,  text , k=5):
    index_name = "semantic-index"
    query_vector = model.encode(text)
    search_body = {
        "size": k,
        "query": {
            "knn": {
                "image_embedding": {
                    "vector": query_vector.tolist(),
                    "k": k,
                }
            }
        },
    }

    result = []
    # {
    #     "image" : 1,
    #     "video" : "video1.mp4"
    #     "score" : 1
    # }
    response = client.search(index=index_name, body=search_body)
    response = response["hits"]["hits"]

    for hit in response:
        image = hit["_source"]["image"]
        video = hit["_source"]["video"]
        score = hit["_score"]
        result.append({"image" : image , "video" : video , "score" : score})

    return result


result_query = query(client=client , model=model , text="A drowning prevention drill is taking place. Flooding occurs, and soldiers are deployed to schools to clean up the floodwaters" , k=5)
print(result_query)

def convert_to_s3_key(video, image):
    #video L03_V025.npy => L03_V025
    # image 18 => 018.jpg
    # image 1 => 001.jpg
    video = video.split(".")[0]
    image = str(image).zfill(3) + ".jpg"
    return video + "/" + image

fig = plt.figure(figsize=(20, 10))
for res in result_query:
    fig.add_subplot(1, 5, result_query.index(res) + 1)
    img = s3.get_object(Bucket=BUCKET_NAME, Key=f"{FOLDER_NAME}/{convert_to_s3_key(res['video'],res['image'])}")["Body"].read()
    img = Image.open(io.BytesIO(img))
    plt.imshow(img)

plt.show()