Upload gallery_review.py

gallery_review.py

@@ -0,0 +1,316 @@
+import os
+import torch
+import numpy as np
+import lightning.pytorch as pl
+import gradio as gr
+import imageio
+import random
+import matplotlib.pyplot as plt
+import cv2
+
+from torch.utils.data import Dataset, DataLoader
+
+from PIL import Image
+from matplotlib import cm
+
+from minimal_script import EmbeddingNetworkSmall, closest_interval, down_to_1k
+from sklearn.cluster import AgglomerativeClustering
+from sklearn.manifold import TSNE
+from sklearn.neighbors import KDTree
+
+
+class PLModule(pl.LightningModule):
+    def __init__(self):
+        super().__init__()
+        self.save_hyperparameters()
+        self.network = EmbeddingNetworkSmall()
+
+    def forward(self, x):
+        return self.network(x)
+
+    def predict_step(self, batch, batch_idx, dataloader_idx=0):
+        outputs = self.forward(batch[0])
+        return outputs, batch[1]
+
+
+class PredictDataset(Dataset):
+    def __init__(self, data_dir, sample=None):
+        self.image_paths = []
+        extensions = ('jpg', 'jpeg', 'png', 'tif', 'webp')
+        for fname in sorted(os.listdir(data_dir)):
+            if any(fname.lower().endswith(ext) for ext in extensions):
+                self.image_paths.append(os.path.join(data_dir, fname))
+        if sample:
+            self.image_paths = random.sample(self.image_paths, sample)
+
+    def __len__(self):
+        return len(self.image_paths)
+
+    def __getitem__(self, idx):
+        path = self.image_paths[idx]
+        image = imageio.v3.imread(path).copy()
+        image = torch.from_numpy(image).permute(2, 0, 1)
+        processed = closest_interval(down_to_1k(image, 1024))
+        processed = 2*(processed/255)-1
+        return processed.detach(), path
+
+
+def explore_embedding_space(embeddings, image_paths, model):
+    """
+    Create an interface for exploring N-dimensional image embeddings
+
+    Args:
+        embeddings: NumPy array of shape [B, N]
+        image_paths: List of B image file paths
+    """
+    # Validate inputs
+    assert len(embeddings) == len(image_paths), "Mismatch between embeddings and image paths"
+    assert embeddings.ndim == 2, "Embeddings should be 2-dimensional"
+
+    # Precompute min/max for each dimension
+    min_vals = embeddings.min(axis=0)
+    max_vals = embeddings.max(axis=0)
+    ranges = max_vals - min_vals
+
+    # Build KDTree for efficient nearest neighbor search
+    tree = KDTree(embeddings)
+
+    # Create initial point (mean of embeddings)
+    initial_point = embeddings.mean(axis=0).tolist()
+
+    # Create slider components for each dimension
+    sliders = []
+    for i in range(embeddings.shape[1]):
+        slider = gr.Slider(
+            float(min_vals[i]),
+            float(max_vals[i]),
+            value=float(initial_point[i]),
+            step=float(ranges[i]) / 100,
+            label=f"Dimension {i + 1}"
+        )
+        sliders.append(slider)
+
+    def compute_gradient_heatmap(image_path):
+        """Compute gradient heatmap for an image"""
+        # Load and preprocess image
+        img = imageio.v3.imread(image_path).copy()
+        img = torch.from_numpy(img).permute(2, 0, 1)
+        img_tensor = closest_interval(down_to_1k(img, 1024)).unsqueeze(0)
+        img_tensor = 2*(img_tensor/255)-1
+        img_tensor.requires_grad_(True)
+
+        # Move to GPU if available
+        device = 'cuda' if torch.cuda.is_available() else 'cpu'
+        img_tensor = img_tensor.to(device)
+
+        # Compute embedding and gradient
+        with torch.enable_grad():
+            embd = model(img_tensor)
+            norm = embd.norm(p=2, dim=1).sum()
+            grad = torch.autograd.grad(norm, img_tensor, retain_graph=False)[0]
+
+        # Compute gradient magnitude
+        grad_mag = grad.squeeze(0).norm(dim=0).detach().cpu().numpy()
+
+        # Normalize and apply colormap
+        grad_min, grad_max = grad_mag.min(), grad_mag.max()
+        if grad_max > grad_min:
+            grad_norm = (grad_mag - grad_min) / (grad_max - grad_min)
+        else:
+            grad_norm = grad_mag * 0  # Handle uniform case
+
+        heatmap = cm.jet(grad_norm)[..., :3]  # Use jet colormap
+        return heatmap
+
+    def overlay_heatmap(original_img, heatmap, alpha=0.6):
+        """Overlay heatmap on original image"""
+        # Resize heatmap to match original image
+        heatmap_img = Image.fromarray((heatmap * 255).astype(np.uint8))
+        heatmap_img = heatmap_img.resize(original_img.size)
+
+        # Convert original to RGBA and heatmap to RGBA
+        #original_rgba = original_img.convert("RGBA")
+        #heatmap_rgba = heatmap_img.convert("RGBA")
+
+        # Blend images
+        blended = Image.blend(original_img, heatmap_img, alpha)
+        return blended
+
+    def get_overlay_image(image_path):
+        """Get image with gradient overlay"""
+        img = Image.open(image_path).convert('RGB')
+        heatmap = compute_gradient_heatmap(image_path)
+        return overlay_heatmap(img, heatmap)
+        #return img
+
+    def add_caption_to_image(image, caption):
+        """Add text caption to the bottom of an image"""
+        # Convert to OpenCV format
+        if isinstance(image, Image.Image):
+            img = np.array(image)
+        else:
+            img = image.copy()
+
+        # Add black bar at bottom
+        bar_height = 30
+        img = cv2.copyMakeBorder(img, 0, bar_height, 0, 0, cv2.BORDER_CONSTANT, value=[0, 0, 0])
+
+        # Add white text
+        font = cv2.FONT_HERSHEY_SIMPLEX
+        text_size = cv2.getTextSize(caption, font, 0.5, 1)[0]
+        text_x = (img.shape[1] - text_size[0]) // 2
+        text_y = img.shape[0] - 10
+        cv2.putText(img, caption, (text_x, text_y), font, 0.5, (255, 255, 255), 1)
+
+        return Image.fromarray(img)
+
+    # Function to find nearby images
+    def find_nearby_images(*point):
+        point = np.array(point).reshape(1, -1)
+        distances, indices = tree.query(point, k=8)
+        indices = indices[0]
+        distances = distances[0]
+
+        # Get paths and create overlay images
+        paths = [image_paths[i] for i in indices]
+        images_with_gradients = [get_overlay_image(p) for p in paths]
+
+        # Create images with baked-in captions
+        final_images = []
+        for img, dist in zip(images_with_gradients, distances):
+            caption = f"Dist: {dist:.2f}"
+            final_img = add_caption_to_image(img, caption)
+            final_images.append(final_img)
+
+        warning = ""
+        if distances[0] > 5.0:  # Warn if nearest image is far
+            warning = "⚠️ Nearest image is far (distance={:.2f}). Consider adjusting sliders.".format(distances[0])
+
+        return final_images, warning
+
+    # Build interface
+    with gr.Blocks() as demo:
+        gr.Markdown("## N-Dimensional Embedding Space Explorer")
+        gr.Markdown("Adjust sliders to navigate. Images show gradient of embedding norm w.r.t. input.")
+
+        # Warning output
+        warning = gr.Textbox(label="Status", interactive=False)
+
+        # Gallery for images
+        gallery = gr.Gallery(
+            label="Nearest Images (Distance Ordered)",
+            columns=4,
+            object_fit="contain",
+            height="auto",
+            show_label=True,
+        )
+
+        # Create sliders in a compact row
+        with gr.Row():
+            for slider in sliders:
+                slider.render()
+
+        # Connect slider changes to update function
+        for slider in sliders:
+            slider.change(
+                find_nearby_images,
+                inputs=sliders,
+                outputs=[gallery, warning]
+            )
+
+        # Initial trigger
+        demo.load(
+            find_nearby_images,
+            inputs=sliders,
+            outputs=[gallery, warning]
+        )
+
+    return demo
+
+
+
+def generate_embeddings(image_folder, mode, model):
+    predict_dataset = PredictDataset(image_folder, 1000)
+    predict_loader = DataLoader(predict_dataset, batch_size=1, num_workers=5, pin_memory=True)
+    trainer = pl.Trainer(accelerator="gpu", logger=False, enable_checkpointing=False)
+    predictions_0 = trainer.predict(model, predict_loader)
+    predictions = torch.cat([pred[0] for pred in predictions_0], dim=0).numpy()
+    paths = []
+    for pred in predictions_0:
+        for i in pred[1]:
+            paths.append(i)
+    if mode == 'Grouping':
+        labels = cluster_embeddings(predictions)
+
+        row_norms = np.linalg.norm(predictions, axis=1)
+        average_norms = np.mean(np.abs(predictions), axis=0)
+        plt.figure(figsize=(8, 5))
+        plt.bar(range(predictions.shape[1]), average_norms, color='skyblue')
+        plt.xlabel('Feature Index (C)')
+        plt.ylabel('Average Norm')
+        plt.title(f'Average Norm for Each Feature (Column)')
+        plt.xticks(range(predictions.shape[1]))
+        plt.show()
+
+        plt.figure(figsize=(8, 6))
+        tsne = TSNE(n_components=2, random_state=42)
+        reduced_data = tsne.fit_transform(predictions)
+        plt.scatter(reduced_data[:, 0], reduced_data[:, 1], c=row_norms, cmap='viridis', s=50, edgecolor='k', label="Data Points")
+        plt.colorbar(label='Norm Value')
+        plt.xlabel('Feature 1')
+        plt.ylabel('Feature 2')
+        plt.title(f'Scatter Plot of Data Points and Average Norm')
+        plt.legend()
+        plt.grid(True)
+        plt.axis('equal')
+        plt.show()
+
+        # List unique clusters
+        unique_clusters = np.unique(labels)
+        # Gradio UI
+        with gr.Blocks() as demo:
+            gr.Markdown("## Explore Image Clusters by Style")
+
+            # Dropdown for selecting a cluster
+            cluster_selector = gr.Dropdown(choices=unique_clusters.tolist(), label="Select Cluster to Explore")
+
+            # Gallery to display images
+            image_gallery = gr.Gallery(label="Sample Images from Selected Cluster")
+
+
+            # Gradio Interface for Cluster Exploration
+            def explore_clusters(cluster_idx):
+                # Find images that belong to the selected cluster
+                cluster_images = [paths[i] for i in range(len(labels)) if labels[i] == cluster_idx]
+                # Load and return images
+                images = [Image.open(img_path) for img_path in cluster_images[:50]]  # Show a sample of 50 images
+                return images
+
+            # Update function for the gallery
+            cluster_selector.change(fn=explore_clusters, inputs=cluster_selector, outputs=image_gallery)
+
+        demo.launch()
+    elif mode == 'Explore':
+        demo = explore_embedding_space(predictions, paths, model.to('cuda'))
+        demo.launch()
+
+
+# Apply Agglomerative Clustering
+def cluster_embeddings(predictions, distance_threshold=6.0):
+    agg_clustering = AgglomerativeClustering(
+        n_clusters=None,
+        distance_threshold=distance_threshold,
+        linkage='ward'
+    )
+    labels = agg_clustering.fit_predict(predictions)
+    return labels
+
+
+
+if __name__ == '__main__':
+    folder = 'Enter Images folder name here'
+    #folder = 'images_for_style_embedding'
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    model = PLModule.load_from_checkpoint('Final_8.ckpt')
+    # 'Grouping' or 'Explore'
+    generate_embeddings(folder, 'Grouping', model)