|
import PIL.Image |
|
import os |
|
import google.generativeai as genai |
|
from prompt import prompt |
|
import google.ai.generativelanguage as glm |
|
import json |
|
|
|
|
|
os.environ['GOOGLE_API_KEY'] = "AIzaSyDAG_Xl66vh4ceY81UXe3vrdwP6wIAkpBs" |
|
genai.configure(api_key=os.environ['GOOGLE_API_KEY']) |
|
|
|
|
|
|
|
def get_image_data(image): |
|
|
|
try: |
|
vision_model = genai.GenerativeModel('gemini-1.5-flash') |
|
except Exception as e: |
|
print(f"Error initializing vision model: {e}") |
|
exit() |
|
|
|
|
|
try: |
|
response = vision_model.generate_content([prompt, image], |
|
generation_config=genai.types.GenerationConfig( |
|
candidate_count=1, |
|
stop_sequences=['.'], |
|
max_output_tokens=200, |
|
top_p=0.7, |
|
top_k=4, |
|
temperature=0.7, |
|
) |
|
) |
|
|
|
|
|
response_text = response.text.strip().lower() |
|
category = response_text.split('category:')[1].split('\n')[0].strip() if 'category:' in response_text else 'Unknown' |
|
type_name = response_text.split('type:')[1].strip() if 'type:' in response_text else 'Unknown' |
|
|
|
|
|
result = { |
|
"Category": category.capitalize(), |
|
"Type": type_name.capitalize() |
|
} |
|
|
|
|
|
return(result) |
|
|
|
except Exception as e: |
|
print(f"Error generating content: {e}") |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|