Upload 2 files

app.py

@@ -0,0 +1,194 @@
+import torch
+from PIL import Image
+from transformers import CLIPProcessor, CLIPModel
+from pathlib import Path
+from torch.utils.data import Dataset, DataLoader
+import os
+import numpy as np
+from numpy.linalg import norm
+import matplotlib.pyplot as plt
+import gradio as gr
+
+# Cast 1 -------------------------------------------------------------------------
+
+def get_clip_embeddings(input_data, input_type='text'):
+    # Load the CLIP model and processor
+    model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32")
+    processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32")
+
+    # Prepare the input based on the type
+    if input_type == 'text':
+        inputs = processor(text=input_data, return_tensors="pt", padding=True, truncation=True)
+    elif input_type == 'image':
+        if isinstance(input_data, str):
+            image = Image.open(input_data)
+        elif isinstance(input_data, Image.Image):
+            image = input_data
+        else:
+            raise ValueError("For image input, provide either a file path or a PIL Image object")
+        inputs = processor(images=image, return_tensors="pt")
+    else:
+        raise ValueError("Invalid input_type. Choose 'text' or 'image'")
+
+    # Get the embeddings
+    with torch.no_grad():
+        if input_type == 'text':
+            embeddings = model.get_text_features(**inputs)
+        else:
+            embeddings = model.get_image_features(**inputs)
+
+    return embeddings.numpy()
+
+# Cast 2 -------------------------------------------------------------------------
+
+class ImageDataset(Dataset):
+    def __init__(self, image_dir, processor):
+        self.image_paths = [os.path.join(image_dir, f) for f in os.listdir(image_dir) if f.endswith(('.png', '.jpg', '.jpeg'))]
+        self.processor = processor
+
+    def __len__(self):
+        return len(self.image_paths)
+
+    def __getitem__(self, idx):
+        image = Image.open(self.image_paths[idx])
+        return self.processor(images=image, return_tensors="pt")['pixel_values'][0]
+
+def get_clip_embeddings_batch(image_dir, batch_size=32, device='cuda'):
+    # Load the CLIP model and processor
+    model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32").to(device)
+    processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32")
+
+    # Create dataset and dataloader
+    dataset = ImageDataset(image_dir, processor)
+    dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=False, num_workers=4)
+
+    all_embeddings = []
+
+    model.eval()
+    with torch.no_grad():
+        for batch in dataloader:
+            batch = batch.to(device)
+            embeddings = model.get_image_features(pixel_values=batch)
+            all_embeddings.append(embeddings.cpu().numpy())
+
+    return np.concatenate(all_embeddings)
+
+# Cast 3 -------------------------------------------------------------------------
+
+# Funkcia na výpočet cosinovej similarity
+def cosine_similarity(x, y):
+    return np.dot(x, y) / (norm(x) * norm(y))
+
+# Funkcia na nájdenie indexov obrázkov
+def maxCS_indices(text_input, embeddings):
+    text_embedding = get_clip_embeddings(text_input, input_type='text')
+    
+    x = text_embedding
+    Y = embeddings
+    
+    # print("Text embedding shape:", x.shape)
+    # print("Embeddings shape:", Y.shape)
+
+    # Vypočítaj cosinovú similaritu pre každý riadok matice Y
+    cosine_similarities = np.array([cosine_similarity(x, y) for y in Y])
+
+    # Získaj indexy štyroch vektorov s najväčšou cosinovou similaritou
+    maxCS_indices = np.argsort(cosine_similarities, axis = 0)[-4:]
+
+    # Výsledné vektory
+    least_similar_vectors = Y[maxCS_indices]
+
+    # print("Indexy vektorov s najmenšou cosinovou similaritou:", smallest_indices)
+    # print("Vektory s najmenšou cosinovou similaritou:\n", least_similar_vectors)
+
+    return(maxCS_indices)
+
+# Cast 4 -------------------------------------------------------------------------
+
+def which_images(images_folder, indices):
+    # Získání všech názvů obrázků ve složce
+    image_filenames = [f for f in os.listdir(images_folder) if f.endswith(('.jpg', '.png'))]
+
+    # Vytvoření numpy array z názvů obrázků
+    image_names_array = np.array(image_filenames)
+
+    # Vytvorenie vektora
+    image_names = (image_names_array[indices])
+    # print(image_names_array[smallest_indices])
+
+    # Transformácia z poľa na vektor
+    image_names_final = image_names.flatten()
+    # print(image_names_final)
+
+    return(image_names_final)
+
+# Cast 5 -------------------------------------------------------------------------
+
+def display_images(folder_path, image_names):
+    # Zložka s obrázkami
+    folder = Path(folder_path)
+
+    # Inicializuj subplots pre 2x2 grid (4 obrázky)
+    fig, axes = plt.subplots(1, 4, figsize=(20, 5))
+
+    # Prejdi cez všetky zadané obrázky a vykresli ich
+    for ax, img_name in zip(axes, image_names):
+        # Načítaj obrázok
+        img_path = folder / img_name
+        img = Image.open(img_path)
+
+        # Vykresli obrázok na subplot
+        ax.imshow(img)
+        ax.set_title(img_name)  # Nastav názov obrázka ako titulok
+        ax.axis('off')          # Skry výstup osí
+
+    # Zobraz obrázky
+    plt.show()
+
+# Cast 6 -------------------------------------------------------------------------
+
+# Nastavenie parametrov pre funkciu process_input
+images_folder = "kotlarska2/Trains"
+device = "cuda" if torch.cuda.is_available() else "cpu"
+embeddings = get_clip_embeddings_batch(images_folder, 32, device)
+
+# Hlavná funkcia na spracovanie vstupu a zobrazenie obrázkov
+def process_input(text_input):
+    our_indinces = maxCS_indices(text_input, embeddings)
+    our_images = which_images(images_folder, our_indinces)
+    return display_our_images(images_folder, our_images)
+
+# Funkcia na zobrazenie obrázkov
+def display_our_images(folder_path, image_names):
+    # Zložka s obrázkami
+    folder = Path(folder_path)
+
+    # Inicializuj subplots pre 2x2 grid (4 obrázky)
+    fig, axes = plt.subplots(1, 4, figsize=(20, 5))
+
+    # Prejdi cez všetky zadané obrázky a vykresli ich
+    for ax, img_name in zip(axes, image_names):
+        # Načítaj obrázok
+        img_path = folder / img_name
+        img = Image.open(img_path)
+
+        # Vykresli obrázok na subplot
+        ax.imshow(img)
+        ax.set_title(img_name)  # Nastav názov obrázka ako titulok
+        ax.axis('off')          # Skry výstup osí
+
+    # Ulož obrázok do súboru
+    plt.tight_layout()
+    plt.savefig('output_images.png')
+    plt.close()
+    
+    return 'output_images.png'
+
+# Nastav a spusti Gradio rozhranie
+iface = gr.Interface(
+    fn=process_input, 
+    inputs="text", 
+    outputs="image", 
+    title="Image Similarity", 
+    description="Zadaj text a zobrazia sa 4 najpodobnejšie obrázky z našej databázy SUV vozidiel.")
+iface.launch(share=True)

requirements.txt

@@ -0,0 +1,11 @@
+numpy==1.26.4
+scipy==1.11.4
+scikit-learn==1.3.2
+fastai==2.7.17
+gradio==4.44.1
+timm==1.0.9
+torch==2.2.1
+torchvision==0.17.1
+transformers==4.29.0
+Pillow==9.4.0
+setuptools