# -*- coding: utf-8 -*-
"""Demo 1: Natural Language Visual Search Of Television News Using OpenAI's CLIP

Automatically generated by Colaboratory.

Original file is located at
    https://colab.research.google.com/drive/1igvHa6oOtGtNGyRV677KStmWmH2oL078
"""

!pip install sentence-transformers

import sentence_transformers
from sentence_transformers import SentenceTransformer, util
from PIL import Image
import glob
import torch
import pickle
import zipfile
from IPython.display import display
from IPython.display import Image as IPImage
import os
from tqdm.autonotebook import tqdm
torch.set_num_threads(4)

import re

#First, we load the respective CLIP model
model = SentenceTransformer('clip-ViT-L-14')

#set our showid and UNIX start time and download and unzip...
    #set UNIX start time by opening the show in the Visual Explorer (https://api.gdeltproject.org/api/v2/tvv/tvv?id=RUSSIA1_20221024_083000_60_minut) 
    #and clicking on first thumbnail (top-left) and copy-pasting the "play" parameter to be "starttime" below...

    starttime = 1666600200
    showid = 'RUSSIA1_20221024_083000_60_minut'

    if not os.path.exists(showid + '.zip'):
        util.http_get('https://storage.googleapis.com/data.gdeltproject.org/gdeltv3/iatv/visualexplorer/' + showid + '.zip', showid + '.zip')
        
    #Extract all images
    with zipfile.ZipFile(showid + '.zip', 'r') as zf:
        for member in tqdm(zf.infolist(), desc='Extracting'):
           zf.extract(member)

#compute the embeddings
img_names = list(glob.glob(showid + '/*.jpg'))
print("Images:", len(img_names))
img_emb = model.encode([Image.open(filepath) for filepath in img_names], batch_size=128, convert_to_tensor=True, show_progress_bar=True)

def search(query, k=5):
    # First, we encode the query (which can either be an image or a text string)
    query_emb = model.encode([query], convert_to_tensor=True, show_progress_bar=False)
    
    # Then, we use the util.semantic_search function, which computes the cosine-similarity
    # between the query embedding and all image embeddings.
    # It then returns the top_k highest ranked images, which we output
    hits = util.semantic_search(query_emb, img_emb, top_k=k)[0]
    
    print("Query:")
    display(query)
    for hit in hits:
        print(img_names[hit['corpus_id']])
        m = re.search('(\d+)\.jpg',img_names[hit['corpus_id']])
        print('https://api.gdeltproject.org/api/v2/tvv/tvv?id=' + showid + '&play=' + str(((int(m.group(1)) + 0) * 4) + starttime) )
        display(IPImage(os.path.join(img_names[hit['corpus_id']]), width=200))

search("nuclear")

search("ukrainian flag")

search("person in front of flag")

search("person in front of bookcase")

search("bookcase")

search("rally")

search("angry")

search("happy")

search("sad")

search("military")

search("helicopter landing")

search("helicopter flying")

search("blue and green")

search("gold and blue")

search("winter jacket")

search("two people talking")

search("panel discussion")

search("warship")

search("zelensky")

search("putin")

search("ukraine")

search("white house")

search("books")

search("podium")

search("aerial imagery")

search("newspaper")

search("map")

search("gun")

search("combat")