bgRemoval-right.tsx 8.9 KB
import { Button, Flex, message } from 'antd';
import { DownloadOutlined, PictureOutlined } from '@ant-design/icons';
import { useEffect, useState } from 'react';
import type { FileProps } from '@/app/types/zuotang';
import { ApiPath } from '@/app/constant';
import { useAccessStore } from '@/app/store';
import type { LocalData } from '@/app/types/zuotang';

// 错误消息映射函数
const getErrorMessage = (state: number): string => {
    const errorMap: { [key: number]: string } = {
        [-8]: '处理超时,最长处理时间30秒',
        [-7]: '无效图片文件(可能已损坏或格式错误)',
        [-5]: '图片大小超过15MB限制',
        [-3]: '服务器下载图片失败,请检查URL有效性',
        [-2]: '处理结果上传失败',
        [-1]: '任务处理失败'
    };
    return errorMap[state] || `未知错误(状态码:${state})`;
};

// 图片URL转Blob方法
const urlToBlob = async (url: string): Promise<Blob> => {
    const response = await fetch(url);
    if (!response.ok) throw new Error('图片加载失败');
    return await response.blob();
};

// 通用轮询函数
const useTaskPoller = () => {
    const pollTask = async (
        taskId: string,
        endpoint: string,
        onSuccess: (data: any) => void,
        onError: (message: string) => void
    ) => {
        const pollInterval = 1000;
        const maxAttempts = 60;
        let attempts = 0;
        const intervalId = setInterval(async () => {
            try {
                attempts++;
                if (attempts > maxAttempts) {
                    clearInterval(intervalId);
                    onError('处理超时,请稍后重试');
                    return;
                }
                const result = await fetch(`${ApiPath.ZuoTang}/${endpoint}/${taskId}`);
                if (!result.ok) {
                    const errorData = await result.json();
                    throw new Error(errorData.message || '状态查询失败');
                }
                const taskResult = await result.json();
                // 根据 state 字段处理状态
                switch (taskResult.data?.state) {
                    case 1: // 任务成功
                        clearInterval(intervalId);
                        onSuccess(taskResult.data);
                        break;
                    case 0: // 队列中
                    case 2: // 准备中
                    case 3: // 等待中
                    case 4: // 处理中
                        // 保持轮询不做操作
                        break;
                    default: // 处理错误状态
                        clearInterval(intervalId);
                        if (taskResult.data?.state < 0) { // 所有负数状态均为错误
                            onError(getErrorMessage(taskResult.data.state));
                        } else {
                            onError('未知任务状态');
                        }
                }
            } catch (error) {
                clearInterval(intervalId);
                onError(error instanceof Error ? error.message : '请求异常');
            }
        }, pollInterval);

        return () => clearInterval(intervalId);
    };

    return { pollTask };
};

// 日期处理工具
const useDateUtils = () => {
    const getFormattedToday = (): string => {
        const today = new Date();
        return today.toISOString().split('T')[0].replace(/-/g, '');
    };

    const isToday = (dateStr: string): boolean => {
        return dateStr === getFormattedToday();
    };

    return { getFormattedToday, isToday };
};

// 本地存储管理
const useLocalStorage = () => {
    const { getFormattedToday, isToday } = useDateUtils();

    const getLocalData = (accessCode: string): LocalData => {
        try {
            const data = localStorage.getItem(accessCode);
            if (!data) return defaultLocalData();

            const parsed = JSON.parse(data);
            if (!isToday(parsed.date)) return defaultLocalData();

            return {
                date: parsed.date || getFormattedToday(),
                maxDailyUses: parsed.maxDailyUses || 'first'
            };
        } catch (e) {
            return defaultLocalData();
        }
    };

    const updateLocalUsage = (accessCode: string, maxDailyUses: number) => {
        const saveData: LocalData = {
            date: getFormattedToday(),
            maxDailyUses: maxDailyUses.toString()
        };
        localStorage.setItem(accessCode, JSON.stringify(saveData));
    };

    const defaultLocalData = (): LocalData => ({
        date: getFormattedToday(),
        maxDailyUses: 'first'
    });

    return { getLocalData, updateLocalUsage };
};

export function BgRemovalRight(props: FileProps) {
    const { previewUrl, setPreviewUrl, fileData,setFileData, isLoading, setIsLoading } = props;
    const accessStore = useAccessStore();
    const { pollTask } = useTaskPoller();
    const { updateLocalUsage, getLocalData } = useLocalStorage();

    const handleApiRequest = async (endpoint: string) => {
        if (!previewUrl) return message.error('请选择图片');
        if (!accessStore.accessCode) return message.error('请先输入访问密码');

        try {
            const formData = new FormData();
            const localData = getLocalData(accessStore.accessCode); // 获取本地数据
            formData.append("accessCode", accessStore.accessCode);
            formData.append("localData", JSON.stringify(localData)); // 序列化后添加
            formData.append("image_file", fileData as Blob);

            const res = await fetch(`${ApiPath.ZuoTang}/${endpoint}`, {
                method: 'POST',
                body: formData
            });
            if (!res.ok) {
                const errorData = await res.json();
                throw new Error(errorData.message || '请求失败');
            }
            const responseData = await res.json();
            if (responseData.status >= 400) {
                if (responseData.status === 429) {
                    updateLocalUsage(accessStore.accessCode, 0);
                }
                throw new Error(responseData.message);
            }
            return responseData;
        } finally {

        }
    };

    const handleProcessImage = async (endpoint: string) => {
        setIsLoading(true);
        try {
            const responseData = await handleApiRequest(endpoint);
            updateLocalUsage(accessStore.accessCode, responseData.maxDailyUses);

            pollTask(
                responseData.data.task_id,
                endpoint,
                async (data) => {
                    try {
                        // 获取新图片的Blob
                        const newBlob = await urlToBlob(data.image || data.image_1);
                        // 创建新的对象URL
                        const newUrl = URL.createObjectURL(newBlob);
                        // 同步更新所有相关状态
                        setPreviewUrl(newUrl);
                        setFileData(newBlob);
                        message.success('处理成功');
                    } catch (error) {
                        message.error('结果图片加载失败');
                    } finally {
                        setIsLoading(false);
                    }
                },
                (errorMsg) => {
                    message.error(errorMsg);
                    setIsLoading(false);
                }
            );
        } catch (error) {
            // 异常时关闭loading
            setIsLoading(false);
            message.error(error instanceof Error ? error.message : '处理失败');
        }
    };
    const handleClick = () => handleProcessImage('visual/segmentation');
    const generateBackground = () => handleProcessImage('visual/r-background');

    const handleDownload = () => {
        if (!previewUrl) return message.error("请先完成图片处理");
        const link = document.createElement('a');
        link.href = previewUrl;
        link.download = `processed-${Date.now()}.png`;
        document.body.appendChild(link);
        link.click();
        document.body.removeChild(link);
    };

    return (
        <Flex vertical justify="center" align="center" gap="middle" style={{ height: '100%' }}>
            <Button
                icon={<PictureOutlined />}
                size="large"
                onClick={handleClick}
                loading={isLoading}
            >
                立即抠图
            </Button>
            <Button
                icon={<DownloadOutlined />}
                size="large"
                onClick={handleDownload}
                disabled={!previewUrl}
            >
                下载图片
            </Button>
            <Button
                icon={<PictureOutlined />}
                size="large"
                onClick={generateBackground}
                loading={isLoading}
            >
                生成背景
            </Button>
        </Flex>
    );
}