app.py

import torch
from transformers import BlipProcessor, BlipForConditionalGeneration
from PIL import Image
import re
import gradio as gr
import os

class ALTTextGenerator:
    def __init__(self):
        """Initialize the BLIP model for fast inference"""
        self.device = "cuda" if torch.cuda.is_available() else "cpu"
        print(f"Using device: {self.device}")
        
        # Use BLIP-1 model - much smaller and memory efficient (~2GB)
        model_name = "Salesforce/blip-image-captioning-base"
        
        # Load model with optimizations for memory efficiency
        self.processor = BlipProcessor.from_pretrained(model_name)
        self.model = BlipForConditionalGeneration.from_pretrained(
            model_name,
            torch_dtype=torch.float16 if self.device == "cuda" else torch.float32,
            device_map="auto" if self.device == "cuda" else None,
            low_cpu_mem_usage=True
        )
        
        if self.device == "cuda":
            self.model.half()  # Use half precision for speed
        
        self.model.eval()
        
        # Prefixes to remove from generated text
        self.prefixes_to_remove = [
            "a photo of ", "an image of ", "a picture of ", "a photograph of ",
            "the image shows ", "this image shows ", "the photo shows ",
            "this photo shows ", "shown in the image is ", "visible in the image is ",
            "the image depicts ", "this image depicts ", "a view of ", "an image showing ",
            "in the image, ", "the picture shows ", "this picture shows ", "there is ",
            "there are ", "we can see ", "you can see "
        ]
    
    def clean_description(self, text):
        """Clean and format the generated description"""
        # Remove any potential conversation formatting
        text = text.strip()
        
        # Convert to lowercase for prefix matching
        text_lower = text.lower().strip()
        
        # Remove common prefixes
        for prefix in self.prefixes_to_remove:
            if text_lower.startswith(prefix):
                text = text[len(prefix):].strip()
                break
        
        # Capitalize first letter
        if text:
            text = text[0].upper() + text[1:] if len(text) > 1 else text.upper()
        
        # Remove multiple spaces and clean up
        text = re.sub(r'\s+', ' ', text).strip()
        
        # Remove trailing period if present
        text = text.rstrip('.')
        
        # Split into sentences and take the first complete sentence if too long
        sentences = re.split(r'[.!?]+', text)
        if sentences:
            text = sentences[0].strip()
        
        return text
    
    def enforce_length_constraints(self, text, min_chars=25, max_chars=125):
        """Ensure text meets length requirements"""
        if len(text) < min_chars:
            # If too short, try to expand with more detail
            expanded = text + " with clear visual details"
            return expanded if len(expanded) <= max_chars else text
        
        if len(text) > max_chars:
            # Truncate at word boundary
            truncated = text[:max_chars]
            last_space = truncated.rfind(' ')
            if last_space > max_chars * 0.75:  # If we can find a reasonable word boundary
                text = truncated[:last_space]
            else:
                text = truncated
        
        return text.strip()
    
    def generate_alt_text(self, image_path):
        """Generate ALT text for the given image"""
        try:
            # Load and process image
            if isinstance(image_path, str):
                image = Image.open(image_path).convert('RGB')
            else:
                image = image_path.convert('RGB')
            
            # Process image
            inputs = self.processor(
                images=image, 
                return_tensors="pt"
            ).to(self.device)
            
            # Generate with optimized parameters for speed and quality
            with torch.no_grad():
                generated_ids = self.model.generate(
                    **inputs,
                    do_sample=True,
                    temperature=0.7,
                    max_length=50,
                    min_length=15,
                    top_p=0.9,
                    num_beams=4,
                    early_stopping=True,
                    repetition_penalty=1.1
                )
            
            # Decode the generated text
            generated_text = self.processor.batch_decode(
                generated_ids, 
                skip_special_tokens=True
            )[0]
            
            # Clean and format the description
            cleaned_text = self.clean_description(generated_text)
            
            # Enforce length constraints
            final_text = self.enforce_length_constraints(cleaned_text)
            
            return final_text
            
        except Exception as e:
            return f"Error processing image: {str(e)}"
    
    def generate_detailed_alt_text(self, image_path):
        """Generate more detailed ALT text with conditional prompting"""
        try:
            if isinstance(image_path, str):
                image = Image.open(image_path).convert('RGB')
            else:
                image = image_path.convert('RGB')
            
            # BLIP supports conditional generation with text prompts
            text_prompt = "describe this image in detail:"
            
            inputs = self.processor(
                images=image,
                text=text_prompt,
                return_tensors="pt"
            ).to(self.device)
            
            with torch.no_grad():
                generated_ids = self.model.generate(
                    **inputs,
                    do_sample=True,
                    temperature=0.8,
                    max_length=60,
                    min_length=20,
                    top_p=0.95,
                    num_beams=5,
                    repetition_penalty=1.2
                )
            
            generated_text = self.processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
            
            # Remove the prompt from the generated text if present
            if text_prompt in generated_text.lower():
                generated_text = generated_text.lower().replace(text_prompt, "").strip()
            
            cleaned_text = self.clean_description(generated_text)
            final_text = self.enforce_length_constraints(cleaned_text)
            
            return final_text
            
        except Exception as e:
            return f"Error processing image: {str(e)}"
    
    def batch_generate(self, image_paths):
        """Generate ALT text for multiple images"""
        results = []
        for img_path in image_paths:
            alt_text = self.generate_alt_text(img_path)
            results.append({
                'image_path': img_path,
                'alt_text': alt_text,
                'char_count': len(alt_text)
            })
        return results

def create_gradio_interface():
    """Create a Gradio interface for the ALT text generator"""
    generator = ALTTextGenerator()
    
    def process_image(image, detail_level):
        if detail_level == "Standard":
            alt_text = generator.generate_alt_text(image)
        else:
            alt_text = generator.generate_detailed_alt_text(image)
        
        char_count = len(alt_text)
        
        # Provide feedback on length
        if char_count < 25:
            status = f"⚠ Too short ({char_count} chars) - Consider adding more detail"
        elif char_count > 125:
            status = f"⚠ Too long ({char_count} chars) - Consider shortening"
        else:
            status = f"✓ Good length ({char_count} chars)"
        
        return alt_text, f"Character count: {char_count}/125\n{status}"
    
    # Create Gradio interface
    interface = gr.Interface(
        fn=process_image,
        inputs=[
            gr.Image(type="pil", label="Upload Image"),
            gr.Radio(
                choices=["Standard", "Detailed"], 
                value="Standard", 
                label="Detail Level",
                info="Standard: Quick generation, Detailed: More comprehensive descriptions"
            )
        ],
        outputs=[
            gr.Textbox(label="Generated ALT Text", lines=4, show_copy_button=True),
            gr.Textbox(label="Length Analysis", lines=2)
        ],
        title="🖼️ ALT Text Generator",
        
        examples=None,
        cache_examples=False,
        theme=gr.themes.Soft()
    )
    
    return interface

def main():
    """Main function to run the ALT text generator"""
    print("Initializing BLIP ALT Text Generator...")
    print("This model is optimized for Hugging Face Spaces memory constraints...")
    
    # Use with Gradio interface
    print("Starting Gradio interface...")
    interface = create_gradio_interface()
    interface.launch(
        share=True,
        server_name="0.0.0.0",
        server_port=7860
    )

if __name__ == "__main__":
    main()