Image_text_generation.py

import os
import io
import tempfile
from datetime import datetime
import PIL.Image
from google import genai
from google.genai import types
import httpx

genai_client = genai.Client(api_key=os.environ["GOOGLE_API_KEY"])

class Image_text_Generator:
    def __init__(self):
        pass

    def generate_image_with_gemini(self, prompt):
        """
        使用 Gemini 模型生成圖片。

        參數:
        prompt (str): 用於生成圖片的提示詞。

        返回:
        bytes: 生成的圖片的二進位資料，如果生成失敗則返回 None。
        """
        response = genai_client.models.generate_content(
            model="gemini-2.0-flash",
            contents=prompt,
            config=types.GenerateContentConfig(response_modalities=['Text', 'Image'])
        )

        for part in response.candidates[0].content.parts:
            if part.text is not None:
                print(part.text)
            elif part.inline_data is not None:
                return part.inline_data.data
        return None

    def upload_image_to_tmp(self, user_id):
        """
        返回:
        str: 上傳後的圖片 URL，如果上傳失敗則返回 None。
        """
        try:
            # 將二進位資料轉換為 PIL Image
            image = PIL.Image.open(io.BytesIO(image_binary))

            # 建立暫存檔案以便上傳
            image.save('tmp/'+user_id+'.png', format='PNG')

            # 返回圖片的連結
            return 'tmp/'+user_id+'.png'
        except Exception as e:
            print(f"圖片上傳失敗: {e}")
            return None

    def generate_txt_with_gemini(self, prompt,doc_url):
        doc_data = httpx.get(doc_url).content

        response = client.models.generate_content(
            model="gemini-2.0-flash",
            contents=[
                types.Part.from_bytes(data=doc_data, mime_type='application/pdf',),prompt]
        )
        
        return response.text