nano_api.py

import io
import requests
import json
import base64
import re
import os
import os.path as osp
import datetime
from typing import Optional, Tuple
from PIL import Image
from io import BytesIO
from tqdm import tqdm
# from api_class.utils import base64_to_image,encode_pil_to_base64
def encode_pil_to_base64(image_pil):
    # 将PIL图像编码为base64字符串
    buffered = BytesIO()
    image_pil.save(buffered, format="PNG")
    img_bytes = buffered.getvalue()
    img_base64 = base64.b64encode(img_bytes).decode('utf-8')
    mime_type = "image/png"
    return img_base64, mime_type

def base64_to_image(base64_str):
    # 将base64字符串解码为PIL图像
    img_bytes = base64.b64decode(base64_str)
    image_pil = Image.open(BytesIO(img_bytes))
    return image_pil
class GeminiImageGenerator:
    def __init__(self,  api_url: str = "https://api.apiyi.com/v1beta/models/gemini-3-pro-image-preview:generateContent"):
        self.api_key = "sk-MC5B3H948s5YhiVN591f578fC74a4eC484659cC6005bB603"
        self.api_url = api_url
        self.api_url_compre="https://api.apiyi.com/v1/chat/completions"
        self.headers = {
            "Content-Type": "application/json",
            "Authorization": f"Bearer {self.api_key}"
        }
        self.SUPPORTED_ASPECT_RATIOS = [
        "21:9", "16:9", "4:3", "3:2", "1:1",
        "9:16", "3:4", "2:3", "5:4", "4:5"]
        
        self.SUPPORTED_RESOLUTION=["1K","2K","4K"]

    def nano_imageEditing(self,data_dict):
        required_keys = {"prompt","image_list","ratio","resolution"}
        assert required_keys <= data_dict.keys(), \
            f"缺少必要字段，必须包含: {required_keys}，实际提供: {list(data_dict.keys())}"
        ratio=data_dict["ratio"]
        resolution=data_dict["resolution"]
        assert ratio in self.SUPPORTED_ASPECT_RATIOS, f"不支持的比例,支持比例为{self.SUPPORTED_ASPECT_RATIOS}"
        assert resolution in self.SUPPORTED_RESOLUTION,f"不支持的分辨率,支持分辨率为{self.SUPPORTED_RESOLUTION}"
        prompt=data_dict["prompt"]
        img_payload=[]
        for image_pil in data_dict["image_list"]:
            image_pil = Image.open(image_pil)
            image_base64, mime_type=encode_pil_to_base64(image_pil)
            img_payload.append({
                            "inline_data": {
                                "mime_type": mime_type,
                                "data": image_base64
                            }
                        })
        try:
            payload = {
                "contents": [{
                    "parts": [{"text": prompt}]+img_payload
                }]
            }

            if ratio:
                payload["generationConfig"] = {
                    "responseModalities": ["IMAGE"],
                    "imageConfig": {
                        "aspectRatio": ratio,
                        "image_size": resolution
                    }
                }

            print("📡 发送请求到 Gemini API...")
            # 发送非流式请求
            response = requests.post(
                self.api_url,
                headers=self.headers,
                json=payload,
                timeout=120
            )
            
            print("✅ API请求成功，正在解析响应...")
            
            # 解析非流式JSON响应
            try:
                result = response.json()
                print("✅ 成功解析JSON响应")
            except json.JSONDecodeError as e:
                return False, f"JSON解析失败: {str(e)}",None
            
            if "candidates" not in result or len(result["candidates"]) == 0:
                return False, "未找到图片数据", None

            candidate = result["candidates"][0]
            if "content" not in candidate or "parts" not in candidate["content"]:
                return False, "响应格式错误",None

            parts = candidate["content"]["parts"]
            output_image_data = None

            for part in parts:
                if "inlineData" in part and "data" in part["inlineData"]:
                    output_image_data = part["inlineData"]["data"]
                    break

            if not output_image_data:
                return False, "未找到图片数据",None
            
            try:            
                pil_img=base64_to_image(output_image_data)
                return pil_img
            except Exception as e:
                raise ValueError(f"图片加载失败: {e}")
    
                
        except requests.exceptions.Timeout:
            raise RuntimeError("请求超时（300秒）")
        except requests.exceptions.ConnectionError as e:
            raise RuntimeError(f"连接错误: {str(e)}")
        except Exception as e:
            raise ValueError(f"未知错误: {str(e)}")



    def nano_text2image(self, data_dict) :
        required_keys = {"prompt", "ratio","resolution"}
        assert required_keys <= data_dict.keys(), \
            f"缺少必要字段，必须包含: {required_keys}，实际提供: {list(data_dict.keys())}"
        ratio=data_dict["ratio"]
        assert ratio in self.SUPPORTED_ASPECT_RATIOS, f"不支持的比例,支持比例为{self.SUPPORTED_ASPECT_RATIOS}"
        
        prompt="帮我生成图片,图片提示词如下: "+data_dict["prompt"]
        resolution=data_dict["resolution"]
        print("🚀 开始生成图片...")
        print(f"提示词: {prompt}")

        try:
            # 构建请求数据
            payload = {
                "contents": [{
                    "parts": [{"text": prompt}]
                }]
            }

            if ratio:
                payload["generationConfig"] = {
                    "responseModalities": ["IMAGE"],
                    "imageConfig": {
                        "aspectRatio": ratio,
                        "image_size": resolution
                    }
                }

            print("📡 发送请求到 Gemini API...")
            # 发送非流式请求

            response = requests.post(
                self.api_url,
                headers=self.headers,
                json=payload,
                timeout=120
            )
            
            if response.status_code != 200:
                error_msg = f"API请求失败，状态码: {response.status_code}"
                try:
                    error_detail = response.json()
                    error_msg += f", 错误详情: {error_detail}"
                except:
                    error_msg += f", 响应内容: {response.text[:500]}"
                return False, error_msg,None
            
            print("✅ API请求成功，正在解析响应...")
            
            # 解析非流式JSON响应
            try:
                result = response.json()
                print("✅ 成功解析JSON响应")
            except json.JSONDecodeError as e:
                return False, f"JSON解析失败: {str(e)}",None
            
            #提取图片数据
            if "candidates" not in result or len(result["candidates"]) == 0:
                return False, "未找到图片数据",None

            candidate = result["candidates"][0]
            if "content" not in candidate or "parts" not in candidate["content"]:
                return False, "响应格式错误",None

            parts = candidate["content"]["parts"]
            image_data = None

            for part in parts:
                if "inlineData" in part and "data" in part["inlineData"]:
                    image_data = part["inlineData"]["data"]
                    break

            if not image_data:
                return False, "未找到图片数据",None
            try:            
                pil_img=base64_to_image(image_data)
                return pil_img
            except Exception as e:
                raise ValueError(f"图片加载失败: {e}")
                
        except requests.exceptions.Timeout:
            raise RuntimeError("请求超时（300秒）")
        except requests.exceptions.ConnectionError as e:
            raise RuntimeError(f"连接错误: {str(e)}")
        except Exception as e:
            raise ValueError(f"未知错误: {str(e)}")

    def _extract_image_from_base64(self,content: str) -> Tuple[bool, Optional[Image.Image], str]:
        """
        高效提取base64图片并返回PIL Image对象
        
        Args:
            content: 包含图片数据的内容
            
        Returns:
            Tuple[是否成功, PIL Image对象(或None), 消息]
        """
        try:
            print(f"📄 内容预览（前200字符）: {content[:200]}")
            
            # 匹配 base64 图片数据
            base64_pattern = r'data:image/([^;]+);base64,([A-Za-z0-9+/=]+)'
            match = re.search(base64_pattern, content)
            
            if not match:
                print('⚠️  未找到base64图片数据')
                raise ValueError("No image founded!")
            
            image_format = match.group(1)
            b64_data = match.group(2)
            
            print(f'🎨 图像格式: {image_format}')
            print(f'📏 Base64数据长度: {len(b64_data)} 字符')
            
            # 解码 base64
            image_data = base64.b64decode(b64_data)
            
            if len(image_data) < 100:
                return False, None, "解码后的图片数据太小，可能无效"
            
            # 使用 PIL 读取图像
            image = Image.open(io.BytesIO(image_data))
            print(f'🖼️ 图片加载成功，尺寸: {image.size}, 模式: {image.mode}')
            
            return True, image, f"成功提取图像 ({image_format})"
            
        except Exception as e:
            return False, None, f"处理图片时发生错误: {str(e)}"

    def nano_image_comprehension(self, data_dict,prompt):
        required_keys = {"image"}
        assert required_keys <= data_dict.keys(), \
            f"缺少必要字段，必须包含: {required_keys}，实际提供: {list(data_dict.keys())}"
        im=data_dict["image"]
        im = Image.open(im)
        im_base64,_=encode_pil_to_base64(im)
        headers = {
                "Authorization": f"Bearer {self.api_key}",
                "Content-Type": "application/json"
            }
        # Prepare payload - generated image first, then ground truth
        payload = {
            "model": "gemini-2.5-flash",
            "messages": [
                {
                    "role": "user",
                    "content": [
                        {"type": "text", "text": prompt},
                        {
                            "type": "image_url",
                            "image_url": {"url": im_base64}
                        },
                    ]
                }
            ],
            "max_tokens": 5000
        }

        # Make API call
        response = requests.post(self.api_url_compre, headers=headers, json=payload, timeout=600)
        response.raise_for_status()

            
        # Extract response content
        response_data = response.json()
        content = response_data['choices'][0]['message']['content']
        
        print(content)
        return content

if __name__=="__main__":
    g = GeminiImageGenerator()
    # 描述性提示词（根据需要修改）
    prompt = """You are a concise storyboard narrator focused on core scene and composition description. Based on the 1 vertically stitched image containing multiple storyboards (identified as Storyboard 1 to Storyboard N in top-to-bottom order, N = actual number), output ONLY a simple story background and concise composition descriptions for each shot. Strictly follow the JSON format below, with each "Image Composition" limited to ~100 words:
{"Simple Story Background": "1-2 sentences summarizing the basic story context (e.g., 'A girl searches for her lost cat in a suburban neighborhood on a sunny afternoon')","Storyboard_List": [{"Shot Number": 1,"Scene": "Specific location (e.g., front yard of a cottage, forest trail, downtown café)","Image Composition": "Concise description of characters (appearance, posture), key props, framing (shot type: close-up/medium/long/wide), lighting, and core visual elements (max 100 words)","Emotional Tone": "Brief atmosphere (e.g., warm, tense, peaceful)"},{"Shot Number": N,"Scene": "Same as above","Image Composition": "Same as above (max 100 words)","Emotional Tone": "Same as above"}]}
Requirements
All fields are mandatory; no redundant content.
"Image Composition" focuses only on critical visual information (characters, framing, key props, lighting) – no excessive details.
Strictly match the number/order of storyboards in the image (top-to-bottom numbering).
JSON format must be error-free, ready for direct use.
No extra text outside the JSON structure."""

    INPUT_DIR = "dataset/spotlight_sketch_cat/GT"
    OUTPUT_DIR = "dataset/spotlight_sketch_cat"
    RATIO = "16:9"
    # 对输入文件夹内的图片进行排序处理

    output_path = os.path.join(OUTPUT_DIR, "spotlight_nano_comprehension_1203.txt")
    os.makedirs(OUTPUT_DIR, exist_ok=True)

    input_files = sorted(
        fname for fname in os.listdir(INPUT_DIR)
        if os.path.isfile(os.path.join(INPUT_DIR, fname))
    )

    for idx, fname in tqdm(enumerate(input_files), total=len(input_files)):
        src_path = os.path.join(INPUT_DIR, fname)
        # 调用nano_image_comprehension
        # try:
        # pil_in = Image.open(src_path).convert("RGB")

        result = g.nano_image_comprehension({
            "image": src_path,
        },prompt)
        base_name = os.path.splitext(fname)[0]
        with open(output_path, "a", encoding="utf-8") as f:
            result = result.replace("\n", "")
            result = result.replace("```", "")
            result = result.replace("json", "")
            result = result.replace('"Simple Story Background"', f'"Image_Name": "{base_name}", "Simple Story Background"')

            f.write(result.strip("\n") + "\n")
        # except Exception as e:
        #     print(f"处理文件 {fname} 时出错: {e}")