Spaces:
Build error
Build error
| import gradio as gr | |
| import torch | |
| import json | |
| import os | |
| from PIL import Image | |
| import torchvision.transforms as transforms | |
| from transformers import CLIPModel, CLIPTokenizer, CLIPProcessor | |
| import time | |
| # Global variables for model, tokenizer, and processor | |
| model = None | |
| tokenizer = None | |
| processor = None | |
| def load_model(progress=gr.Progress()): | |
| global model, tokenizer, processor | |
| if model is None or tokenizer is None: | |
| progress(0, desc="Starting model load...") | |
| try: | |
| # Load directly from Hugging Face | |
| model_name = "openai/clip-vit-base-patch32" | |
| progress(0.2, desc=f"Loading tokenizer from {model_name}...") | |
| tokenizer = CLIPTokenizer.from_pretrained(model_name) | |
| progress(0.4, desc=f"Loading model from {model_name}...") | |
| model = CLIPModel.from_pretrained(model_name) | |
| progress(0.6, desc="Loading processor...") | |
| processor = CLIPProcessor.from_pretrained(model_name) | |
| # Load quantized weights and convert back to float32 | |
| progress(0.7, desc="Loading quantized weights...") | |
| try: | |
| quantized_weights = torch.load("pytorch_model_quantized.bin", map_location="cpu", weights_only=True) | |
| # Dequantize weights | |
| progress(0.8, desc="Dequantizing weights...") | |
| dequantized_weights = {} | |
| for key, tensor in quantized_weights.items(): | |
| if tensor.dtype == torch.int8: | |
| dequantized_weights[key] = tensor.to(torch.float32) | |
| else: | |
| dequantized_weights[key] = tensor | |
| # Load weights | |
| model.load_state_dict(dequantized_weights) | |
| print("✅ Custom model weights loaded successfully!") | |
| except Exception as weight_error: | |
| print(f"⚠️ Could not load custom weights: {str(weight_error)}") | |
| print("⚠️ Using default model weights instead.") | |
| model.eval() | |
| progress(0.9, desc="Finalizing setup...") | |
| print("✅ Model and tokenizer loaded successfully!") | |
| progress(1.0, desc="Ready!") | |
| except Exception as e: | |
| print(f"❌ Error loading model: {str(e)}") | |
| raise e | |
| return model, tokenizer, processor | |
| # Load the analysis templates | |
| with open("cifar10_image_analysis.json", "r") as f: | |
| analysis_config = json.load(f) | |
| # Set up the image transform | |
| transform = transforms.Compose([ | |
| transforms.Resize((224, 224)), | |
| transforms.ToTensor(), | |
| transforms.Normalize((0.48145466, 0.4578275, 0.40821073), | |
| (0.26862954, 0.26130258, 0.27577711)) | |
| ]) | |
| def analyze_image(image_path, progress=gr.Progress()): | |
| try: | |
| # Load the model with progress bar | |
| progress(0.1, desc="Loading model...") | |
| model, tokenizer, processor = load_model(progress) | |
| # Load and transform the image | |
| progress(0.3, desc="Processing image...") | |
| image = Image.open(image_path) | |
| image_tensor = transform(image).unsqueeze(0) | |
| # Get image features | |
| progress(0.5, desc="Analyzing image...") | |
| with torch.no_grad(): | |
| image_features = model.get_image_features(pixel_values=image_tensor) | |
| results = [] | |
| # Analyze using each template | |
| for i, response in enumerate(analysis_config['responses']): | |
| progress(0.6 + (i * 0.08), desc=f"Answering question {i+1}/5...") | |
| question = response['question'] | |
| template = response['template'] | |
| # Get the class name from the image path | |
| image_name = os.path.basename(image_path) | |
| class_name = image_name.split('.')[0] # Remove .png extension | |
| # Get class-specific information | |
| class_info = analysis_config['class_info'][class_name] | |
| # Tokenize the question | |
| text_inputs = tokenizer( | |
| question, | |
| padding=True, | |
| return_tensors="pt", | |
| max_length=77, | |
| truncation=True | |
| ) | |
| # Get text features | |
| text_features = model.get_text_features( | |
| input_ids=text_inputs.input_ids, | |
| attention_mask=text_inputs.attention_mask | |
| ) | |
| # Calculate similarity | |
| similarity = torch.nn.functional.cosine_similarity( | |
| image_features, text_features, dim=1 | |
| ) | |
| # Format the answer | |
| answer = template.format( | |
| class_name=class_name, | |
| object_type=class_info['object_type'] | |
| ) | |
| # Add more detailed answer for the description question | |
| if question == "Describe this image in detail.": | |
| answer += f" {class_info['description']}" | |
| results.append({ | |
| 'question': question, | |
| 'answer': answer, | |
| 'confidence': f"{similarity.item():.2f}" | |
| }) | |
| progress(0.95, desc="Formatting results...") | |
| # Format the output | |
| output = "✨ Analysis Results ✨\n\n" | |
| for result in results: | |
| output += f"❓ Q: {result['question']}\n" | |
| output += f"💡 A: {result['answer']}\n" | |
| output += f"📊 Confidence: {result['confidence']}\n" | |
| output += "=" * 80 + "\n\n" | |
| progress(1.0, desc="Done!") | |
| return output | |
| except Exception as e: | |
| error_msg = f"❌ Error analyzing image: {str(e)}" | |
| print(error_msg) | |
| return error_msg | |
| def create_interface(): | |
| # Get the list of available images | |
| image_files = [f for f in os.listdir("images") if f.endswith(".png")] | |
| image_files.sort() # Sort alphabetically | |
| # Create the interface | |
| with gr.Blocks(theme=gr.themes.Soft(), title="CIFAR-10 Image Analysis with SmolVLM2") as demo: | |
| # Add custom JavaScript to replace any static text | |
| gr.HTML(""" | |
| <script> | |
| // Wait for the DOM to fully load | |
| window.addEventListener('DOMContentLoaded', (event) => { | |
| // Replace all occurrences of "CLIP model" with "SmolVLM2" | |
| document.body.innerHTML = document.body.innerHTML.replace(/CLIP model/g, "SmolVLM2"); | |
| // Add a small delay to catch any elements that might load after | |
| setTimeout(() => { | |
| document.body.innerHTML = document.body.innerHTML.replace(/CLIP model/g, "SmolVLM2"); | |
| }, 1000); | |
| }); | |
| </script> | |
| """) | |
| gr.Markdown( | |
| """ | |
| # 🖼️ CIFAR-10 Image Analysis | |
| Select an image from the dropdown menu and click 'Analyze Image' to get detailed analysis using our fine-tuned SmolVLM2. | |
| """ | |
| ) | |
| # Static text and components | |
| interface_title = "Select an image from the dropdown menu and click 'Analyze Image' to get detailed analysis using our fine-tuned SmolVLM2." | |
| with gr.Row(): | |
| with gr.Column(scale=1): | |
| image_dropdown = gr.Dropdown( | |
| choices=image_files, | |
| label="Select Image", | |
| value=image_files[0] if image_files else None, | |
| container=True | |
| ) | |
| analyze_button = gr.Button("🔍 Analyze Image", variant="primary") | |
| with gr.Column(scale=1): | |
| image_display = gr.Image(label="Selected Image", container=True) | |
| output_text = gr.Textbox( | |
| label="Analysis Results", | |
| lines=20, | |
| container=True | |
| ) | |
| gr.Markdown( | |
| """ | |
| ### 📝 Analysis includes: | |
| - Main object identification | |
| - Color analysis | |
| - Object type (natural/man-made) | |
| - Detailed description | |
| - Background analysis | |
| """ | |
| ) | |
| def update_image(image_name): | |
| return f"images/{image_name}" | |
| def run_analysis(image_name): | |
| return analyze_image(f"images/{image_name}") | |
| image_dropdown.change( | |
| fn=update_image, | |
| inputs=image_dropdown, | |
| outputs=image_display | |
| ) | |
| analyze_button.click( | |
| fn=run_analysis, | |
| inputs=image_dropdown, | |
| outputs=output_text | |
| ) | |
| return demo | |
| # Create and launch the interface | |
| if __name__ == "__main__": | |
| demo = create_interface() | |
| demo.queue() # Enable queuing for better handling of multiple requests | |
| demo.launch() | |