Spaces:
Sleeping
Sleeping
File size: 3,396 Bytes
314c39a fdce92b 3af7c5f fdce92b | 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 105 106 107 108 | import streamlit as st
from transformers import VisionEncoderDecoderModel, ViTImageProcessor, AutoTokenizer
import torch
from PIL import Image
import requests
import os
# Load the model and tokenizer
model = VisionEncoderDecoderModel.from_pretrained("nlpconnect/vit-gpt2-image-captioning")
feature_extractor = ViTImageProcessor.from_pretrained("nlpconnect/vit-gpt2-image-captioning")
tokenizer = AutoTokenizer.from_pretrained("nlpconnect/vit-gpt2-image-captioning")
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)
# Define the image captioning function
def generate_caption(image):
pixel_values = feature_extractor(images=image, return_tensors="pt").pixel_values
pixel_values = pixel_values.to(device)
output_ids = model.generate(pixel_values, max_length=16, num_beams=4)
caption = tokenizer.batch_decode(output_ids, skip_special_tokens=True)[0]
return caption.strip()
# Function to fetch product information from GROQ API
def fetch_product_info(caption):
# Define your GROQ API endpoint
api_url = "GroqCloud" # Replace with your actual GROQ API URL
# Retrieve the API key from environment variables
api_key = os.getenv('groq_api') # Use environment variable for API key
if not api_key:
st.error("API key not found. Set the 'GROQ_API_KEY' environment variable.")
return None
# Set up the headers with the Bearer token
headers = {"Authorization": f"Bearer {api_key}"}
# Query the API
params = {"query": caption}
try:
response = requests.get(api_url, headers=headers, params=params)
response.raise_for_status() # Raise an exception for HTTP errors
data = response.json()
if data:
product = data.get('products', [])[0] # Assuming the first product is relevant
ingredients = product.get('ingredients', 'N/A')
usage = product.get('usage', 'N/A')
barcode = product.get('barcode', 'N/A')
return ingredients, usage, barcode
else:
st.write("No products found in the response.")
return None
except requests.RequestException as e:
st.error(f"Error fetching data from GROQ API: {e}")
return None
# Streamlit UI
st.title("Image Captioning and Product Information")
st.write("Upload an image to get a caption and related product information.")
# Upload image
uploaded_image = st.file_uploader("Choose an image...", type=["jpg", "png", "jpeg"])
if uploaded_image:
image = Image.open(uploaded_image).convert("RGB")
st.image(image, caption="Uploaded Image", use_column_width=True)
# Generate and display caption
with st.spinner("Generating caption..."):
caption = generate_caption(image)
st.write(f"**Caption:** {caption}")
# Fetch and display additional information from GROQ API
info = fetch_product_info(caption)
if info:
ingredients, usage, barcode = info
st.write("### Additional Information:")
st.write(f"**Ingredients:** {ingredients}")
st.write(f"**Usage:** {usage}")
st.write(f"**Barcoding:** {barcode}")
else:
st.write("No additional information found for this product.")
|