3d-sim-platform / frontend /api /gemini_service.py
prolau's picture
Upload folder using huggingface_hub
c84aa63 verified
Raw
History Blame Contribute Delete
3.01 kB
import google.generativeai as genai
import os
import json
import io
from PIL import Image
api_key = os.getenv("GEMINI_API_KEY", "MOCK_KEY")
if api_key != "MOCK_KEY":
genai.configure(api_key=api_key)
class GeminiService:
def __init__(self):
self.model_name = 'gemini-1.5-flash'
self.knowledge_base = ""
try:
base_dir = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
t4bk_path = os.path.join(base_dir, 't4bk.md')
with open(t4bk_path, 'r', encoding='utf-8') as f:
self.knowledge_base = f.read()
print("Successfully loaded historical knowledge base.")
except Exception as e:
print("Warning: Could not load t4bk.md knowledge base:", e)
async def analyze_spatial_layout(self, text_prompt: str, image_bytes: bytes = None) -> dict:
if api_key == "MOCK_KEY":
print("WARNING: Gemini API Key not found. Using MOCK data.")
return {
"layout": {
"primary_object": text_prompt,
"lighting": "ambient",
"style": "marble",
"camera_angle": "front-center",
"scale": 1.0
}
}
model = genai.GenerativeModel(self.model_name)
prompt = f"""
Act as a spatial intelligence architect and historical environment reconstructor.
You have been provided with the following historical knowledge base from the Malaysian Sejarah curriculum:
=== HISTORICAL KNOWLEDGE BASE ===
{self.knowledge_base}
=================================
Analyze the following text prompt and the attached image (if any), using the historical knowledge above if relevant.
Output a detailed JSON layout describing the 3D structure, lighting, style, and object placement for a 3D Gaussian Splat scene.
Ensure your generated layout reflects historical accuracy where applicable based on the provided textbook context.
Text Prompt: {text_prompt}
"""
contents = [prompt]
if image_bytes:
try:
img = Image.open(io.BytesIO(image_bytes))
contents.append(img)
except Exception as e:
print("Failed to load image with PIL:", e)
try:
print("Calling Gemini API...")
response = await model.generate_content_async(contents)
print("Gemini API Success.")
json_start = response.text.find("{")
json_end = response.text.rfind("}") + 1
if json_start != -1 and json_end != -1:
return json.loads(response.text[json_start:json_end])
return {"raw": response.text}
except Exception as e:
print("Gemini API Error:", e)
return {"error": str(e)}
gemini_client = GeminiService()