Spaces:
Sleeping
Sleeping
File size: 3,174 Bytes
555f48a 523547d 647e467 02b1c71 523547d 20351ba 523547d 555f48a 523547d 555f48a 523547d 555f48a 523547d 555f48a 523547d 555f48a 523547d 555f48a 523547d 555f48a 523547d 555f48a 523547d 555f48a 523547d 555f48a 523547d 555f48a 523547d 555f48a 523547d 555f48a 647e467 523547d 555f48a 647e467 20351ba 523547d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 | 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() |