File size: 3,524 Bytes
13d7bed 5e351cb 13d7bed 5e351cb 13d7bed 5e351cb 13d7bed 5e351cb 13d7bed 5e351cb fcba7cb 5e351cb fcba7cb 5e351cb 13d7bed 5e351cb 1f43f91 5e351cb 13d7bed 5e351cb 13d7bed 5e351cb 13d7bed 5e351cb 13d7bed 5e351cb 13d7bed fcba7cb 6cd7102 fcba7cb 5e351cb 13d7bed b7b543a 13d7bed b7b543a 13d7bed 5e351cb 13d7bed |
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 99 100 101 102 103 104 |
import torch
from PIL import Image
import requests
import openai
from transformers import (Owlv2Processor, Owlv2ForObjectDetection,
AutoProcessor, AutoModelForMaskGeneration,
BlipProcessor, BlipForConditionalGeneration)
import matplotlib.pyplot as plt
import matplotlib.patches as patches
import base64
import io
import numpy as np
import gradio as gr
import json
import os
from dotenv import load_dotenv
# Load environment variables
load_dotenv()
OPENAI_API_KEY = os.getenv('OPENAI_API_KEY')
openai.api_key = OPENAI_API_KEY
def generate_image_caption(image):
device = 'cuda' if torch.cuda.is_available() else 'cpu'
processor = BlipProcessor.from_pretrained('Salesforce/blip-image-captioning-base')
model = BlipForConditionalGeneration.from_pretrained('Salesforce/blip-image-captioning-base').to(device)
inputs = processor(image, return_tensors='pt').to(device)
out = model.generate(**inputs)
caption = processor.decode(out[0], skip_special_tokens=True)
return caption
def analyze_caption(caption):
messages = [
{
"role": "user",
"content": f"""Your task is to determine if the following image description is surprising or not surprising.
Description: "{caption}"
If the description is surprising, determine which element, figure, or object is making it surprising and write it only in one sentence with no more than 6 words; otherwise, write 'NA'.
Also, rate how surprising the image is on a scale of 1-5, where 1 is not surprising at all and 5 is highly surprising.
Provide the response as a JSON with the following structure:
{{
"label": "[surprising OR not surprising]",
"element": "[element]",
"rating": [1-5]
}}
"""
}
]
response = openai.ChatCompletion.create(
model="gpt-4",
messages=messages,
max_tokens=100,
temperature=0.1
)
return response.choices[0].message.content
# The rest of your functions (process_image_detection, show_mask, etc.) remain the same
def process_and_analyze(image):
if image is None:
return None, "Please upload an image first."
if OPENAI_API_KEY is None:
return None, "OpenAI API key not found in environment variables."
try:
# Handle different input types
if isinstance(image, tuple):
image = image[0] # Take the first element if it's a tuple
if isinstance(image, np.ndarray):
image = Image.fromarray(image)
if not isinstance(image, Image.Image):
raise ValueError("Invalid image format")
# Generate caption
caption = generate_image_caption(image)
# Analyze caption
gpt_response = analyze_caption(caption)
response_data = json.loads(gpt_response)
if response_data["label"].lower() == "surprising" and response_data["element"].lower() != "na":
result_buf = process_image_detection(image, response_data["element"], response_data["rating"])
result_image = Image.open(result_buf)
analysis_text = f"Label: {response_data['label']}\nElement: {response_data['element']}\nRating: {response_data['rating']}/5"
return result_image, analysis_text
else:
return image, "Not Surprising"
except Exception as e:
return None, f"Error processing image: {str(e)}"
# Create Gradio interface remains the same
if __name__ == "__main__":
demo = create_interface()
demo.launch() |