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(""" """) 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()