import os import torch import torch.nn.functional as F import gradio as gr import numpy as np from PIL import Image import torchvision.transforms.functional as TF from matplotlib import colormaps from transformers import AutoModel # ---------------------------- # Configuration # ---------------------------- # Define available models DEFAULT_MODEL_ID = "facebook/dinov3-vits16plus-pretrain-lvd1689m" ALT_MODEL_ID = "facebook/dinov3-vith16plus-pretrain-lvd1689m" AVAILABLE_MODELS = [DEFAULT_MODEL_ID, ALT_MODEL_ID] PATCH_SIZE = 16 DEVICE = "cuda" if torch.cuda.is_available() else "cpu" # Normalization constants (standard for ImageNet) IMAGENET_MEAN = (0.485, 0.456, 0.406) IMAGENET_STD = (0.229, 0.224, 0.225) # ---------------------------- # Model Loading (with caching) # ---------------------------- _model_cache = {} _current_model_id = None model = None # global reference def load_model_from_hub(model_id: str): """Loads a DINOv3 model from the Hugging Face Hub.""" print(f"Loading model '{model_id}' from Hugging Face Hub...") try: token = os.environ.get("HF_TOKEN") # optional, for gated models mdl = AutoModel.from_pretrained(model_id, token=token, trust_remote_code=True) mdl.to(DEVICE).eval() print(f"✅ Model '{model_id}' loaded successfully on device: {DEVICE}") return mdl except Exception as e: print(f"❌ Failed to load model '{model_id}': {e}") raise gr.Error( f"Could not load model '{model_id}'. " "If the model is gated, please accept the terms on its Hugging Face page " "and set HF_TOKEN in your environment. " f"Original error: {e}" ) def get_model(model_id: str): """Return a cached model if available, otherwise load and cache it.""" if model_id in _model_cache: return _model_cache[model_id] mdl = load_model_from_hub(model_id) _model_cache[model_id] = mdl return mdl # Load the default model at startup model = get_model(DEFAULT_MODEL_ID) _current_model_id = DEFAULT_MODEL_ID def _ensure_model(model_id: str): """Ensure the global 'model' matches the dropdown selection.""" global model, _current_model_id if model_id != _current_model_id: model = get_model(model_id) _current_model_id = model_id # ---------------------------- # Helper Functions # ---------------------------- def resize_to_grid(img: Image.Image, long_side: int, patch: int) -> torch.Tensor: """Resizes an image to dimensions that are multiples of the patch size.""" w, h = img.size scale = long_side / max(h, w) new_h = max(patch, int(round(h * scale))) new_w = max(patch, int(round(w * scale))) new_h = ((new_h + patch - 1) // patch) * patch new_w = ((new_w + patch - 1) // patch) * patch return TF.to_tensor(TF.resize(img.convert("RGB"), (new_h, new_w))) def colorize(data: np.ndarray, cmap_name: str = 'viridis') -> Image.Image: """Converts a 2D numpy array to a colored PIL image.""" x = data.astype(np.float32) x = (x - x.min()) / (x.max() - x.min() + 1e-8) cmap = colormaps.get_cmap(cmap_name) rgb = (cmap(x)[..., :3] * 255).astype(np.uint8) return Image.fromarray(rgb) def blend(base: Image.Image, heat: Image.Image, alpha: float) -> Image.Image: """Blends a heatmap onto a base image.""" base = base.convert("RGBA") heat = heat.convert("RGBA") return Image.blend(base, heat, alpha=alpha) # ---------------------------- # Core Gradio Function # ---------------------------- @torch.inference_mode() def generate_pca_visuals( image_pil: Image.Image, resolution: int, cmap_name: str, overlay_alpha: float, model_id: str, progress=gr.Progress(track_tqdm=True) ): """Main function to generate PCA visuals.""" _ensure_model(model_id) if model is None: raise gr.Error("DINOv3 model is not available. Check the startup logs.") if image_pil is None: return None, None, "Please upload an image and click Generate.", None, None # 1. Image Preprocessing progress(0.2, desc="Resizing and preprocessing image...") image_tensor = resize_to_grid(image_pil, resolution, PATCH_SIZE) t_norm = TF.normalize(image_tensor, IMAGENET_MEAN, IMAGENET_STD).unsqueeze(0).to(DEVICE) original_processed_image = TF.to_pil_image(image_tensor) _, _, H, W = t_norm.shape Hp, Wp = H // PATCH_SIZE, W // PATCH_SIZE # 2. Feature Extraction progress(0.5, desc="🦖 Extracting features with DINOv3...") outputs = model(t_norm) # The model output includes a [CLS] token AND 4 register tokens. n_special_tokens = 5 patch_embeddings = outputs.last_hidden_state.squeeze(0)[n_special_tokens:, :] # 3. PCA Calculation progress(0.8, desc="🔬 Performing PCA...") X_centered = patch_embeddings.float() - patch_embeddings.float().mean(0, keepdim=True) U, S, V = torch.pca_lowrank(X_centered, q=3, center=False) # Stabilize the signs of the eigenvectors for deterministic output. for i in range(V.shape[1]): max_abs_idx = torch.argmax(torch.abs(V[:, i])) if V[max_abs_idx, i] < 0: V[:, i] *= -1 scores = X_centered @ V[:, :3] # 4. Explained Variance total_variance = (X_centered ** 2).sum() explained_variance = [float((s**2) / total_variance) for s in S] variance_text = ( f"**📊 Explained Variance Ratios:**\n\n" f"- **PC1:** {explained_variance[0]:.2%}\n" f"- **PC2:** {explained_variance[1]:.2%}\n" f"- **PC3:** {explained_variance[2]:.2%}" ) # 5. Create Visualizations pc1_map = scores[:, 0].reshape(Hp, Wp).cpu().numpy() pc1_image_raw = colorize(pc1_map, cmap_name) pc_rgb_map = scores.reshape(Hp, Wp, 3).cpu().numpy() min_vals = pc_rgb_map.reshape(-1, 3).min(axis=0) max_vals = pc_rgb_map.reshape(-1, 3).max(axis=0) pc_rgb_map = (pc_rgb_map - min_vals) / (max_vals - min_vals + 1e-8) pc_rgb_image_raw = Image.fromarray((pc_rgb_map * 255).astype(np.uint8)) target_size = original_processed_image.size pc1_image_smooth = pc1_image_raw.resize(target_size, Image.Resampling.BICUBIC) pc_rgb_image_smooth = pc_rgb_image_raw.resize(target_size, Image.Resampling.BICUBIC) blended_image = blend(original_processed_image, pc1_image_smooth, overlay_alpha) progress(1.0, desc="✅ Done!") return pc1_image_smooth, pc_rgb_image_smooth, variance_text, blended_image, original_processed_image # ---------------------------- # Gradio Interface # ---------------------------- with gr.Blocks(theme=gr.themes.Soft(), title="🦖 DINOv3 PCA Explorer") as demo: gr.Markdown( """ # 🦖 DINOv3 PCA Explorer Upload an image to visualize the principal components of its patch features. This reveals the main axes of semantic variation within the image as understood by the model. """ ) with gr.Row(): with gr.Column(scale=2): input_image = gr.Image(type="pil", label="Upload Image", value="https://images.squarespace-cdn.com/content/v1/607f89e638219e13eee71b1e/1684821560422-SD5V37BAG28BURTLIXUQ/michael-sum-LEpfefQf4rU-unsplash.jpg") with gr.Accordion("⚙️ Visualization Controls", open=True): resolution_slider = gr.Slider( minimum=224, maximum=1024, value=512, step=16, label="Processing Resolution", info="Higher values capture more detail but are slower." ) model_choice = gr.Dropdown( choices=AVAILABLE_MODELS, value=DEFAULT_MODEL_ID, label="Backbone (DINOv3)", info="ViT-S/16+ is smaller & faster; ViT-H/16+ is larger.", ) cmap_dropdown = gr.Dropdown( ['viridis', 'magma', 'inferno', 'plasma', 'cividis', 'jet'], value='viridis', label="Heatmap Colormap" ) alpha_slider = gr.Slider( minimum=0, maximum=1, value=0.5, label="Overlay Opacity" ) run_button = gr.Button("🚀 Generate PCA Visuals", variant="primary") with gr.Column(scale=3): with gr.Tabs(): with gr.TabItem("🖼️ Overlay"): gr.Markdown("Visualize the main heatmap blended with the original image.") output_blended = gr.Image(label="PC1 Heatmap Overlay") output_processed = gr.Image(label="Original Processed Image (at selected resolution)") with gr.TabItem("📊 PCA Outputs"): gr.Markdown("View the raw outputs of the Principal Component Analysis.") output_pc1 = gr.Image(label="PC1 Heatmap (Smoothed)") output_rgb = gr.Image(label="Top 3 PCs as RGB (Smoothed)") output_variance = gr.Markdown(label="Explained Variance") run_button.click( fn=generate_pca_visuals, inputs=[input_image, resolution_slider, cmap_dropdown, alpha_slider, model_choice], outputs=[output_pc1, output_rgb, output_variance, output_blended, output_processed] ) if __name__ == "__main__": demo.launch()