from fastapi import FastAPI  
from pydantic import BaseModel  
import os  

# GCP vertexai 
import vertexai
from vertexai.preview.generative_models import GenerativeModel, Image
import re
import base64


PROJECT_ID = os.getenv('PROJECT_ID')  
REGION = os.getenv('REGION')  


app = FastAPI(  
    title="Image Detection Service",  
    description="Image Detection Service",  
    version="0.1.0"  
)  

def process_base64_image(base64_string):  
    if ',' in base64_string:  
        base64_string = base64_string.split(',')[1]  
    
    try:  
        image_bytes = base64.b64decode(base64_string)  
        return Image.from_bytes(image_bytes)  
    except Exception as e:  
        print(f"图片解码错误: {e}")  
        return None  
    
def clean_string(text):  
    return re.sub(r'[^0-9\-]', '', text)  
    
@app.get("/", summary="default")  
async def root():  
    return {"message": "Hello, World!"}  


@app.post("/detect", summary="识别图片")  

async def detect_image(request_data: dict):  
    base64image = request_data.get('base64image')  
    if base64image == "":
        return {"code": 400, "data":[]}  


    vertexai.init(project=PROJECT_ID, location=REGION)

    image = process_base64_image(base64image) 
    generative_multimodal_model = GenerativeModel("gemini-1.5-pro-002")
    response = generative_multimodal_model.generate_content(['''Based on the text prompt at the top and the sample image above, which of the following 9 images should be selected? Multiple selections are allowed. Please number the positions starting from 1. The format should be as follows:
    ###
    1-3-4''', image])

    detect_position = ""
    if response.text != "":
        detect_position = clean_string(response.text)


    return {"code": 200, "data":detect_position.split("-")}