tw2/qingjian_cs/app/auto_draw.py

from __future__ import annotations

import os
import tempfile
import uuid
from typing import Any

import requests
from .logger import setup_logger

from .config import AUTO_DRAW_TIMEOUT_SECONDS

logger = setup_logger()


async def auto_draw_preview(
    image_url: str,
    customer_id: str,
    requirement: str = "",
) -> dict[str, Any]:
    """
    统一自动作图入口（直调本地链路）：
    1) 下载客户图
    2) 调 Gemini 生成
    3) 上传图绘，返回可外发 URL
    """
    try:
        logger.info("[作图] 开始 customer=%s image=%s", customer_id, image_url)
        from services.service_gemini_stable import GeminiExtractStableService  # type: ignore
        from services.service_tuhui_upload import upload_to_tuhui  # type: ignore
        from .image_quote_analyzer import analyze_image_for_quote, evaluate_generated_image
    except Exception as e:
        logger.error("[作图] 依赖加载失败: %s", e)
        return {"ok": False, "error": f"依赖加载失败:{e}"}

    prompt = requirement.strip() or "按原图做高清修复，保留主体细节，输出清晰可用版本"
    input_path = os.path.join(tempfile.gettempdir(), f"qjcs_in_{uuid.uuid4().hex}.jpg")
    output_path = os.path.join(tempfile.gettempdir(), f"qjcs_out_{uuid.uuid4().hex}.jpg")

    try:
        logger.info("[作图] 识图评估中")
        analysis = await analyze_image_for_quote(
            image_url=image_url,
            customer_text=requirement,
            goods_name="",
        )
        if analysis.get("ok"):
            business_related = str(analysis.get("business_related", "yes")).lower()
            can_do = str(analysis.get("can_do", "partial")).lower()
            if business_related == "no":
                logger.info("[作图] 终止: 非印花/印刷相关")
                return {"ok": False, "error": "非印花/印刷相关，退出"}
            if can_do == "no":
                logger.info("[作图] 终止: 识图判定不可做")
                return {"ok": False, "error": "识图判定不可做，退出"}
            if str(analysis.get("gemini_prompt", "")).strip():
                prompt = str(analysis.get("gemini_prompt", "")).strip()
                logger.info("[作图] 使用识图提示词: %s", prompt[:80])

        logger.info("[作图] 下载原图中")
        headers = {
            "User-Agent": (
                "Mozilla/5.0 (Windows NT 10.0; Win64; x64) "
                "AppleWebKit/537.36 (KHTML, like Gecko) "
                "Chrome/122.0.0.0 Safari/537.36"
            ),
            "Referer": "https://www.taobao.com/",
            "Accept": "image/avif,image/webp,image/apng,image/*,*/*;q=0.8",
        }
        resp = requests.get(image_url, headers=headers, timeout=AUTO_DRAW_TIMEOUT_SECONDS)
        if resp.status_code != 200:
            logger.error("[作图] 原图下载失败: http_%s", resp.status_code)
            return {"ok": False, "error": f"原图下载失败:http_{resp.status_code}"}
        with open(input_path, "wb") as f:
            f.write(resp.content)
        logger.info("[作图] 原图下载完成 size=%s", len(resp.content))

        logger.info("[作图] Gemini 生成中")
        service = GeminiExtractStableService()
        ok_extract, msg_extract, _ = await service.extract_pattern(
            input_path=input_path,
            output_path=output_path,
            custom_prompt=prompt,
            aspect_ratio="1:1",
        )
        if not ok_extract:
            logger.error("[作图] Gemini 生成失败: %s", msg_extract)
            return {"ok": False, "error": f"生成失败:{msg_extract}"}
        if not os.path.exists(output_path):
            logger.error("[作图] Gemini 未产出文件")
            return {"ok": False, "error": "生成失败:未产出文件"}
        logger.info("[作图] Gemini 生成完成")

        logger.info("[作图] 结果评估中")
        review = await evaluate_generated_image(
            original_image_url=image_url,
            generated_image_path=output_path,
            requirement=requirement,
        )
        if not review.get("pass", False):
            logger.info("[作图] 终止: 评估不通过 reason=%s", review.get("reason", ""))
            return {
                "ok": False,
                "need_transfer": True,
                "error": f"评估不通过:{review.get('reason', 'unknown')}",
            }

        logger.info("[作图] 上传图绘中")
        ok_upload, link, _ = await upload_to_tuhui(
            output_path,
            title=f"客户{customer_id[-4:]}-预览图" if customer_id else "预览图",
            description="AI自动作图预览",
            price=1,
        )
        if not ok_upload:
            logger.error("[作图] 图绘上传失败: %s", link)
            return {"ok": False, "error": f"上传失败:{link}"}
        logger.info("[作图] 上传成功 url=%s", link)
        return {"ok": True, "url": str(link)}
    except Exception as e:
        logger.exception("[作图] 异常")
        return {"ok": False, "error": f"作图异常:{e}"}
    finally:
        try:
            if os.path.exists(input_path):
                os.remove(input_path)
        except Exception:
            pass
        try:
            if os.path.exists(output_path):
                os.remove(output_path)
        except Exception:
            pass