import gradio as gr from PIL import Image import torch import os import time import numpy as np # Set CUDA memory configuration to avoid fragmentation os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "max_split_size_mb:128" # Import the models after setting memory configuration from transformers import CLIPProcessor, CLIPModel, BlipProcessor, BlipForConditionalGeneration # Model configuration CLIP_MODEL_ID = "openai/clip-vit-base-patch32" # Fast classification DETAILED_MODEL_ID = "Salesforce/blip-image-captioning-large" # Use original BLIP instead of BLIP-2 USE_GPU = torch.cuda.is_available() # Global variables clip_model = None clip_processor = None detailed_model = None detailed_processor = None def load_clip_model(): """Load the CLIP model for fast classification""" global clip_model, clip_processor # Return if already loaded if clip_model is not None and clip_processor is not None: return True print("Loading CLIP model...") try: # First clear any GPU memory if torch.cuda.is_available(): torch.cuda.empty_cache() # Load processor clip_processor = CLIPProcessor.from_pretrained(CLIP_MODEL_ID) # Load model efficiently if USE_GPU: clip_model = CLIPModel.from_pretrained(CLIP_MODEL_ID).to("cuda") else: clip_model = CLIPModel.from_pretrained(CLIP_MODEL_ID) # Set to evaluation mode clip_model.eval() print("CLIP model loaded successfully!") return True except Exception as e: print(f"Error loading CLIP model: {str(e)}") return False def load_detailed_model(): """Load the BLIP model for detailed image analysis""" global detailed_model, detailed_processor # If already loaded, return if detailed_model is not None and detailed_processor is not None: return True print("Loading BLIP model...") try: # Clear memory first if torch.cuda.is_available(): torch.cuda.empty_cache() # Load processor and model for original BLIP detailed_processor = BlipProcessor.from_pretrained(DETAILED_MODEL_ID) # For older models like BLIP, don't use device_map='auto' or load_in_8bit # Instead, load the model and then move it to the device detailed_model = BlipForConditionalGeneration.from_pretrained( DETAILED_MODEL_ID, torch_dtype=torch.float16 if USE_GPU else torch.float32 ) # Manually move model to GPU if available if USE_GPU: detailed_model = detailed_model.to("cuda") # Set to evaluation mode detailed_model.eval() print("BLIP model loaded successfully!") return True except Exception as e: print(f"Error loading BLIP model: {str(e)}") if "CUDA out of memory" in str(e): print("Not enough GPU memory for the detailed model") return False # Categories for image classification CATEGORIES = [ "a photograph", "a painting", "a drawing", "a digital art", "landscape", "portrait", "cityscape", "animals", "food", "vehicle", "building", "nature", "people", "abstract art", "technology", "interior", "exterior", "night scene", "beach", "mountains", "forest", "water", "flowers", "sports", "a person", "multiple people", "a child", "an elderly person", "a dog", "a cat", "wildlife", "a bird", "a car", "a building", "a presentation slide", "a graph", "a chart", "a diagram", "text document", "a screenshot", "a map", "a table of data", "a scientific figure" ] def get_detailed_analysis(image): """Get detailed analysis from the image using BLIP model""" try: start_time = time.time() # Make sure the model is loaded if not load_detailed_model(): return "Couldn't load detailed analysis model." # Convert numpy array to PIL Image if isinstance(image, np.ndarray): image_pil = Image.fromarray(image).convert('RGB') else: # If somehow it's already a PIL Image image_pil = image.convert('RGB') # Resize the image to improve performance max_size = 600 # Limit to 600px on the longest side width, height = image_pil.size if max(width, height) > max_size: if width > height: new_width = max_size new_height = int(height * (max_size / width)) else: new_height = max_size new_width = int(width * (max_size / height)) image_pil = image_pil.resize((new_width, new_height), Image.LANCZOS) device = "cuda" if USE_GPU else "cpu" # Using an unconditional approach first - this usually works better inputs = detailed_processor(image_pil, return_tensors="pt") if USE_GPU: inputs = {k: v.to(device) for k, v in inputs.items()} with torch.no_grad(): # Get a basic unconditional caption output_ids = detailed_model.generate( **inputs, max_length=50, num_beams=5, do_sample=False, early_stopping=True ) base_description = detailed_processor.decode(output_ids[0], skip_special_tokens=True) # ULTRA-SIMPLE single-word prompts to avoid any echoing analyses = { "text": None, # Text content "chart": None, # Chart analysis "subject": None # Main subject } # Use the base description for context with ultra-simple prompts ultra_simple_prompts = { f"Text in {base_description[:20]}...": "text", f"Charts in {base_description[:20]}...": "chart", f"Subject of {base_description[:20]}...": "subject" } for prompt, analysis_type in ultra_simple_prompts.items(): # Process with prompt inputs = detailed_processor(image_pil, text=prompt, return_tensors="pt") if USE_GPU: inputs = {k: v.to(device) for k, v in inputs.items()} with torch.no_grad(): output_ids = detailed_model.generate( **inputs, max_length=75, num_beams=3, do_sample=True, temperature=0.7, repetition_penalty=1.2, early_stopping=True ) result = detailed_processor.decode(output_ids[0], skip_special_tokens=True) # SUPER AGGRESSIVE cleaning # First, remove anything that looks like a prefix before a colon colon_parts = result.split(":") if len(colon_parts) > 1: # Take everything after the first colon result = ":".join(colon_parts[1:]).strip() # Remove the base description if it appears if base_description in result: result = result.replace(base_description, "").strip() # Remove any part of the prompt for p in ultra_simple_prompts.keys(): if p in result: result = result.replace(p, "").strip() # Remove the first 20 chars of base description if they appear if base_description[:20] in result: result = result.replace(base_description[:20], "").strip() # Remove all common question patterns and filler text remove_patterns = [ "text in", "charts in", "subject of", "in detail", "describe", "this image", "the image", "can you", "do you", "is there", "are there", "i can see", "i see", "there is", "there are", "it looks like", "appears to be", "seems to be", "might be", "could be", "i think", "i believe", "probably", "possibly", "maybe", "it is", "this is", "that is", "these are", "those are", "image shows", "picture shows", "image contains", "picture contains", "in the image", "in this image", "of this image", "from this image", "based on", "according to", "looking at", "from what i can see", "appears to show", "depicts", "represents", "illustrates", "demonstrates", "presents", "displays", "portrays", "reveals", "indicates", "suggests", "we can see", "you can see", "one can see" ] for pattern in remove_patterns: if pattern.lower() in result.lower(): # Find and remove each occurrence lower_result = result.lower() while pattern.lower() in lower_result: idx = lower_result.find(pattern.lower()) if idx >= 0: result = result[:idx] + result[idx+len(pattern):] lower_result = result.lower() # Clean up any punctuation/formatting issues result = result.strip() while result and result[0] in ",.;:?!-": result = result[1:].strip() # Remove "..." if it appears result = result.replace("...", "").strip() # Fix capitalization if result and len(result) > 0: result = result[0].upper() + result[1:] if len(result) > 1 else result[0].upper() analyses[analysis_type] = result # Compose the final output output_text = f"## Detailed Description\n{base_description}\n\n" # Only show relevant sections if analyses['text'] and len(analyses['text']) > 5 and not any(x in analyses['text'].lower() for x in ["no text", "not any text", "can't see", "cannot see", "don't see", "couldn't find"]): output_text += f"## Text Content\n{analyses['text']}\n\n" if analyses['chart'] and len(analyses['chart']) > 5 and not any(x in analyses['chart'].lower() for x in ["no chart", "not any chart", "no graph", "not any graph", "can't see", "cannot see", "don't see", "couldn't find"]): output_text += f"## Chart Analysis\n{analyses['chart']}\n\n" output_text += f"## Main Subject\n{analyses['subject'] or 'Unable to determine main subject.'}" # Clear GPU memory if USE_GPU: torch.cuda.empty_cache() elapsed_time = time.time() - start_time return output_text except Exception as e: print(f"Error in detailed analysis: {str(e)}") # Try to clean up memory in case of error if USE_GPU: torch.cuda.empty_cache() return f"Error in detailed analysis: {str(e)}" def get_clip_classification(image): """Get fast classification using CLIP""" if not load_clip_model(): return [] try: # Process with CLIP inputs = clip_processor( text=CATEGORIES, images=image, return_tensors="pt", padding=True ) # Move to GPU if available if USE_GPU: inputs = {k: v.to("cuda") for k, v in inputs.items()} # Get predictions with torch.inference_mode(): outputs = clip_model(**inputs) # Process results logits_per_image = outputs.logits_per_image probs = logits_per_image.softmax(dim=1) # Get top predictions values, indices = probs[0].topk(8) # Format results return [(CATEGORIES[idx], value.item() * 100) for value, idx in zip(values, indices)] except Exception as e: print(f"Error in CLIP classification: {str(e)}") return [] def process_image(image, get_detailed=False): """Process image with both fast and detailed analysis""" if image is None: return "Please upload an image to analyze." try: # Start timing start_time = time.time() # Preprocess image if hasattr(image, 'mode') and image.mode != 'RGB': image = image.convert('RGB') # Resize for efficiency if max(image.size) > 600: # Smaller max size for better performance ratio = 600 / max(image.size) new_size = (int(image.size[0] * ratio), int(image.size[1] * ratio)) image = image.resize(new_size, Image.LANCZOS) # Get fast classification first categories = get_clip_classification(image) result = "## Image Classification\n" result += "This image appears to contain:\n" for category, confidence in categories: result += f"- {category.title()} ({confidence:.1f}%)\n" # Add detailed analysis if requested if get_detailed: result += "\n## Detailed Analysis\n" detailed_result = get_detailed_analysis(image) result += detailed_result # Add timing information elapsed_time = time.time() - start_time result += f"\n\nAnalysis completed in {elapsed_time:.2f} seconds." # Clean up memory if torch.cuda.is_available(): torch.cuda.empty_cache() return result except Exception as e: print(f"Error: {str(e)}") if torch.cuda.is_available(): torch.cuda.empty_cache() return f"Error processing image: {str(e)}" # Create interface with more options with gr.Blocks(title="Enhanced Image Analyzer") as demo: gr.Markdown("# Enhanced Image Analyzer") gr.Markdown("Upload an image and choose between fast classification or detailed analysis.") with gr.Row(): with gr.Column(): input_image = gr.Image(type="pil", label="Upload an image") detailed_checkbox = gr.Checkbox(label="Get detailed analysis (slower but better quality)", value=False) analyze_btn = gr.Button("Analyze Image", variant="primary") with gr.Column(): output = gr.Markdown(label="Analysis Results") analyze_btn.click( fn=process_image, inputs=[input_image, detailed_checkbox], outputs=output ) # Optional examples if os.path.exists("data_temp"): examples = [os.path.join("data_temp", f) for f in os.listdir("data_temp") if f.endswith(('.png', '.jpg', '.jpeg'))] if examples: gr.Examples(examples=examples, inputs=input_image) if __name__ == "__main__": # Start with clean memory if torch.cuda.is_available(): torch.cuda.empty_cache() demo.launch(server_name="0.0.0.0", server_port=int(os.environ.get("PORT", 7860)))