File size: 5,906 Bytes
1b0da9b |
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 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 |
import os
import json
import base64
from typing import List, Tuple, Dict
import gradio as gr
import httpx
from sqlalchemy import create_engine, text
from dotenv import load_dotenv
import google.generativeai as genai
def get_secret(secret_name, service="", username=""):
try:
from google.colab import userdata
return userdata.get(secret_name)
except:
try:
return os.environ[secret_name]
except:
import keyring
return keyring.get_password(service, username)
# Load environment variables
load_dotenv()
# Database configuration
DB_NAME = "kroyscappingdb"
DB_USER = "read_only"
DB_PASSWORD = get_secret('FASHION_PG_PASS')
DB_HOST = "rc1d-vbh2dw5ha0gpsazk.mdb.yandexcloud.net"
DB_PORT = "6432"
DATABASE_URL = f"postgresql://{DB_USER}:{DB_PASSWORD}@{DB_HOST}:{DB_PORT}/{DB_NAME}"
# Create the SQLAlchemy engine
db_conn = create_engine(DATABASE_URL)
# Configure Gemini API
genai.configure(api_key=get_secret("GEMINI_API_KEY"))
def get_marketplace_and_main_image(id_product_money: str) -> Tuple[str, str]:
"""Get marketplace and main image URL for a product."""
query = text("""
select mp, image as main_image_url
from public.products
where id_product_money = :id_product_money
""")
with db_conn.connect() as connection:
result = connection.execute(query, {"id_product_money": id_product_money}).first()
if result is None:
raise ValueError(f"No product found with id_product_money: {id_product_money}")
return result.mp, result.main_image_url
def get_additional_images(id_product_money: str, marketplace: str) -> List[str]:
"""Get additional images based on marketplace."""
if marketplace == 'lamoda':
query = text("""
select info_chrc->'gallery' as more_images
from public.lamoda_chrc_and_reviews
where id_product_money = :id_product_money
limit 1
""")
with db_conn.connect() as connection:
result = connection.execute(query, {"id_product_money": id_product_money}).first()
if result and result.more_images:
paths = json.loads(result.more_images)
return [f"https://a.lmcdn.ru/product{path}" for path in paths]
elif marketplace == 'wildberries':
query = text("""
select features->'images' as more_images
from public.wb_chrc
where id_product_money = :id_product_money
limit 1
""")
with db_conn.connect() as connection:
result = connection.execute(query, {"id_product_money": id_product_money}).first()
if result and result.more_images:
return json.loads(result.more_images)
return []
def download_and_encode_images(image_urls: List[str]) -> List[Dict]:
"""Download images and convert them to base64 format for Gemini."""
encoded_images = []
with httpx.Client() as client:
for url in image_urls:
try:
response = client.get(url)
response.raise_for_status()
encoded_image = base64.b64encode(response.content).decode('utf-8')
encoded_images.append({
'mime_type': 'image/jpeg', # Assuming JPEG format
'data': encoded_image
})
except Exception as e:
print(f"Error downloading image {url}: {str(e)}")
return encoded_images
def get_gemini_response(model_name: str, encoded_images: List[Dict], prompt: str) -> str:
"""Get response from a Gemini model."""
try:
model = genai.GenerativeModel(model_name)
response = model.generate_content(encoded_images + [prompt])
return response.text
except Exception as e:
return f"Error with {model_name}: {str(e)}"
def process_input(id_product_money: str, prompt: str) -> Tuple[List[str], str, str]:
"""Main processing function."""
try:
# Get marketplace and main image
marketplace, main_image = get_marketplace_and_main_image(id_product_money)
# Get additional images
additional_images = get_additional_images(id_product_money, marketplace)
# Combine all images
all_image_urls = [main_image] + additional_images
# Download and encode images
encoded_images = download_and_encode_images(all_image_urls)
if not encoded_images:
raise ValueError("No images could be downloaded")
# Get responses from both models
gemini_1_5_response = get_gemini_response("gemini-1.5-pro", encoded_images, prompt)
gemini_2_0_response = get_gemini_response("gemini-pro-vision", encoded_images, prompt)
return all_image_urls, gemini_1_5_response, gemini_2_0_response
except Exception as e:
return [], f"Error: {str(e)}", f"Error: {str(e)}"
# Create Gradio interface
with gr.Blocks() as demo:
gr.Markdown("# Product Image Analysis with Gemini Models")
with gr.Row():
id_input = gr.Textbox(label="Product ID (id_product_money)")
prompt_input = gr.Textbox(label="Prompt for VLMs")
submit_btn = gr.Button("Analyze")
with gr.Row():
image_gallery = gr.Gallery(label="Product Images")
with gr.Row():
with gr.Column():
gr.Markdown("### Gemini 1.5 Pro Response")
gemini_1_5_output = gr.Textbox(label="")
with gr.Column():
gr.Markdown("### Gemini Pro Vision Response")
gemini_2_0_output = gr.Textbox(label="")
submit_btn.click(
fn=process_input,
inputs=[id_input, prompt_input],
outputs=[image_gallery, gemini_1_5_output, gemini_2_0_output]
)
if __name__ == "__main__":
demo.launch(share=True)
|