Spaces:
Sleeping
Sleeping
| import torch | |
| import gradio as gr | |
| from PIL import Image | |
| from transformers import ( | |
| BlipProcessor, | |
| BlipForConditionalGeneration, | |
| pipeline | |
| ) | |
| # Select device | |
| device = "cuda" if torch.cuda.is_available() else "cpu" | |
| torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32 | |
| # Load BLIP captioning model directly | |
| processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-base") | |
| caption_model = BlipForConditionalGeneration.from_pretrained( | |
| "Salesforce/blip-image-captioning-base", | |
| torch_dtype=torch_dtype | |
| ).to(device) | |
| # Load image classification model | |
| classifier = pipeline( | |
| task="image-classification", | |
| model="google/vit-base-patch16-224", | |
| device=0 if torch.cuda.is_available() else -1 | |
| ) | |
| print("Models loaded successfully.") | |
| def generate_caption(image): | |
| inputs = processor(images=image, return_tensors="pt").to(device, torch_dtype) | |
| output = caption_model.generate(**inputs, max_new_tokens=30) | |
| caption = processor.decode(output[0], skip_special_tokens=True) | |
| return caption | |
| def build_summary(caption: str, label: str) -> str: | |
| caption = caption.strip() if caption else "No caption available" | |
| label = label.strip() if label else "unknown object" | |
| return ( | |
| f"The captioning model describes the image as: {caption}. " | |
| f"The image classification model identifies the main subject as: {label}. " | |
| f"Taken together, the image appears to focus on this subject or scene." | |
| ) | |
| def analyze_image(image): | |
| try: | |
| if image is None: | |
| return ( | |
| "Please upload an image.", | |
| "No classification available.", | |
| "Please upload an image first." | |
| ) | |
| image = image.convert("RGB") | |
| # Captioning | |
| caption = generate_caption(image) | |
| print("CAPTION RESULT:", caption) | |
| # Classification | |
| class_result = classifier(image) | |
| print("CLASSIFICATION RESULT:", class_result) | |
| if isinstance(class_result, list) and len(class_result) > 0: | |
| top_label = class_result[0].get("label", "Unknown") | |
| top_score = class_result[0].get("score", 0.0) | |
| classification_text = f"{top_label} (confidence: {top_score:.4f})" | |
| else: | |
| top_label = "Unknown" | |
| classification_text = "No classification generated." | |
| summary = build_summary(caption, top_label) | |
| return caption, classification_text, summary | |
| except Exception as e: | |
| print("ERROR:", str(e)) | |
| error_text = f"Error: {str(e)}" | |
| return error_text, error_text, error_text | |
| demo = gr.Interface( | |
| fn=analyze_image, | |
| inputs=gr.Image(type="pil", label="Upload an Image"), | |
| outputs=[ | |
| gr.Textbox(label="Generated Caption"), | |
| gr.Textbox(label="Top Classification"), | |
| gr.Textbox(label="Combined Summary", lines=4) | |
| ], | |
| title="Image Captioning, Classification, and Summary App", | |
| description=( | |
| "Upload an image to generate an automatic caption, predict the main image class, " | |
| "and produce a short combined summary." | |
| ), | |
| ) | |
| if __name__ == "__main__": | |
| demo.launch() |