Spaces:
Running
Running
| 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() |