from sentence_transformers import SentenceTransformer, util
from PIL import Image
import pickle
import os
import gradio as gr
import zipfile
import logging
import datetime

logger = logging.getLogger(__name__)

# Load CLIP model
text_model = SentenceTransformer('clip-ViT-B-32-multilingual-v1')

image_model = SentenceTransformer('clip-ViT-B-32')
image_model.parallel_tokenization = False

img_folder = '.\\photos\\'
if not os.path.exists(img_folder) or len(os.listdir(img_folder)) == 0:
    os.makedirs(img_folder, exist_ok=True)

    photo_filename = 'unsplash-25k-photos.zip'
    if not os.path.exists(photo_filename):  # Download dataset if does not exist
        util.http_get('http://sbert.net/datasets/' +
                      photo_filename, photo_filename)

    # Extract all images
    with zipfile.ZipFile(photo_filename, 'r') as zf:
        for member in zf.infolist():
            zf.extract(member, img_folder)


emb_filename = ".\\unsplash-25k-photos-embeddings.pkl"
if not os.path.exists(emb_filename):
    util.http_get(
        'http://sbert.net/datasets/unsplash-25k-photos-embeddings.pkl', emb_filename)

with open(emb_filename, 'rb') as fIn:
    img_names, img_emb = pickle.load(fIn)


img_folder = ".\\photos\\"
duplicates = util.paraphrase_mining_embeddings(img_emb)


def search_duplicates(threshold=100):
    current_time = datetime.datetime.now()
    logger.warning(f"{current_time} : {threshold}")
    threshold /= 100
    near_duplicates = [entry for entry in duplicates if entry[0] <= threshold]
    images = []
    for score, idx1, idx2 in near_duplicates[0:5]:
        img1_path = os.path.join(img_folder, img_names[idx1])
        img2_path = os.path.join(img_folder, img_names[idx2])
        img1 = Image.open(img1_path)
        img2 = Image.open(img2_path)
        images.append((img1, ""))
        images.append((img2, score))
    return images


description = '''
Приложение для поиска одинаковых изображений - это мощный инструмент, который помогает пользователям выявлять одинаковые изображения и, например, управлять ими.
Это приложение использует передовые алгоритмы компьютерного зрения и машинного обучения для обнаружения похожих или идентичных изображений и представляет их в понятном и организованном виде.
Пользователи могут загружать изображения в приложение и быстро определять похожие изображения либо дупликаты, которые могут существовать в их коллекции. 
Приложение может потенциально предоставить инструменты для управления дубликатами, например, возможность их удаления или объединения.
Приложение для поиска дубликатов изображений идеально подходит для тех, кому необходимо управлять большой коллекцией изображений, например, для фотографов, графических дизайнеров или менеджеров социальных сетей. 
Выявляя дубликаты изображений и управляя ими, пользователи могут сэкономить время и оптимизировать свои рабочие процессы.
Благодаря удобному интерфейсу и мощным алгоритмам приложение для поиска дубликатов изображений является отличным примером того, как ИИ может быть использован для повышения производительности и оптимизации рабочих процессов.
'''

iface_duplicates = gr.Interface(
    title="Поиск похожих изображений - d8a.ai",
    description=description,
    allow_flagging="never",
    fn=search_duplicates,
    inputs=[gr.inputs.Slider(91, 100, step=1, default=100,
                             label="Процент вероятности")],
    outputs=gr.Gallery(
        label="Найденные похожие изображения", elem_id="gallery"
    ).style(grid=[2], height="auto")
)

iface_duplicates.launch()