审查视图

app/components/common/bgRemoval-right.tsx 8.9 KB
202304001 authored
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252
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>
    );
}