forge/frontend/app/image/remove-bg/page.tsx

346 lines
15 KiB
TypeScript

'use client';
import { useState, useEffect } from 'react';
import { useSearchParams } from 'next/navigation';
import { toast } from 'react-hot-toast';
import { Eraser, Download, Sparkles, Trash2, RefreshCw } from 'lucide-react';
import FileUpload from '@/components/FileUpload';
import JobProgress from '@/components/JobProgress';
import { modulesApi, assetsApi, jobsApi } from '@/lib/api';
import { useStore } from '@/lib/store';
const outputFormats = [
{ value: 'png', label: 'PNG (Transparent)' },
{ value: 'webp', label: 'WebP' },
{ value: 'tiff', label: 'TIFF (Clipping Path)' },
];
interface QueueItem {
id: string;
file?: File; // Optional because we might load by assetId from URL
assetId?: string;
originalFileName?: string; // To display when only assetId is known
jobId?: string;
status: 'pending' | 'uploading' | 'processing' | 'completed' | 'error';
resultAssetId?: string;
error?: string;
}
export default function RemoveBackgroundPage() {
const { addJob, updateJob } = useStore();
const searchParams = useSearchParams();
const [queue, setQueue] = useState<QueueItem[]>([]);
const [outputFormat, setOutputFormat] = useState('png');
const [refineMask, setRefineMask] = useState(true);
const [processing, setProcessing] = useState(false);
useEffect(() => {
const urlAssetId = searchParams.get('assetId');
if (urlAssetId) {
// Add as a pending item if not already in queue
setQueue(prev => {
if (prev.some(item => item.assetId === urlAssetId)) return prev;
return [...prev, {
id: Math.random().toString(36).substring(7),
assetId: urlAssetId,
status: 'pending',
originalFileName: 'Asset from URL' // We might not know the name yet, that's fine
}];
});
toast.success('Asset added to queue from URL');
}
}, [searchParams]);
const handleFileUpload = (files: File[]) => {
const newItems: QueueItem[] = files.map(file => ({
id: Math.random().toString(36).substring(7),
file,
originalFileName: file.name,
status: 'pending'
}));
setQueue(prev => [...prev, ...newItems]);
toast.success(`${files.length} images added to queue`);
};
const processItem = async (item: QueueItem) => {
if (item.status === 'completed' || item.status === 'processing') return item;
try {
let assetId = item.assetId;
// 1. Upload if needed
if (!assetId && item.file) {
setQueue(prev => prev.map(i => i.id === item.id ? { ...i, status: 'uploading' } : i));
try {
const uploadRes = await assetsApi.upload(item.file);
assetId = uploadRes.data.id;
} catch (err: any) {
if (err.response?.status === 409) {
const existingAssetId = err.response.data.detail.asset_id;
const shouldOverwrite = window.confirm(
`File "${item.file.name}" already exists. \nClick OK to Overwrite, Cancel to Use Existing.`
);
if (shouldOverwrite) {
const uploadRes = await assetsApi.upload(item.file, undefined, false, true);
assetId = uploadRes.data.id;
} else {
assetId = existingAssetId;
}
} else {
throw err;
}
}
setQueue(prev => prev.map(i => i.id === item.id ? { ...i, assetId, status: 'pending' } : i));
} else if (!assetId && !item.file) {
throw new Error("No file or asset ID provided");
}
// 2. Process
setQueue(prev => prev.map(i => i.id === item.id ? { ...i, status: 'processing' } : i));
// This assumes modulesApi.removeBackground takes { asset_id, output_format, refine_mask }
const response = await modulesApi.removeBackground({
asset_id: assetId!,
output_format: outputFormat,
refine_mask: refineMask,
});
const job = response.data;
addJob({
id: job.id,
module: 'background_removal',
status: job.status,
progress: job.progress,
created_at: job.created_at,
});
// Poll for completion
let currentJob = job;
while (currentJob.status !== 'completed' && currentJob.status !== 'failed') {
await new Promise(resolve => setTimeout(resolve, 2000));
const pollRes = await jobsApi.get(job.id);
if (pollRes?.data) currentJob = pollRes.data;
else break;
}
if (currentJob.status === 'completed' && currentJob.output_asset_ids?.length > 0) {
const resultId = currentJob.output_asset_ids[0];
setQueue(prev => prev.map(i => i.id === item.id ? { ...i, status: 'completed', resultAssetId: resultId } : i));
return { ...item, status: 'completed', resultAssetId: resultId };
} else {
throw new Error(currentJob.error_message || 'Job failed');
}
} catch (err: any) {
console.error(err);
const errMsg = err.message || 'Failed';
setQueue(prev => prev.map(i => i.id === item.id ? { ...i, status: 'error', error: errMsg } : i));
return { ...item, status: 'error', error: errMsg };
}
};
const handleProcessQueue = async () => {
setProcessing(true);
const pending = queue.filter(i => i.status === 'pending' || i.status === 'error');
const limit = 3;
// We process only pending items. Note that if we just uploaded (item.status is 'pending' but with assetId set), it will process.
// If it was 'uploading', it shouldn't happen here because upload is part of processItem if we do it one by one,
// OR we can separates upload step. The current logic puts upload inside processItem which is fine.
for (let i = 0; i < pending.length; i += limit) {
const chunk = pending.slice(i, i + limit);
await Promise.all(chunk.map(item => processItem(item)));
}
setProcessing(false);
toast.success('Queue processing complete');
};
const handleClearQueue = () => {
setQueue([]);
};
const removeItem = (id: string) => {
setQueue(prev => prev.filter(i => i.id !== id));
};
const handleDownload = async (assetId: string) => {
try {
const response = await assetsApi.download(assetId);
// We need the filename. If we haven't fetched the asset details, we might guess or try to get it.
// For now let's hope the browser handles it or we fetch header.
// Actually assetsApi.download returns a blob?
// Better logic: fetch asset details to get filename, then download.
// But here let's just use generic name if we can't get it easily, or maybe we fetch asset first?
// To be safe and quick:
const url = window.URL.createObjectURL(response.data);
const a = document.createElement('a');
a.href = url;
a.download = `removed_bg_${assetId}.png`; // fallback
a.click();
window.URL.revokeObjectURL(url);
} catch (err) {
toast.error("Download failed");
}
};
return (
<div className="max-w-6xl mx-auto space-y-8">
<div className="flex items-center gap-4">
<div className="w-12 h-12 bg-forge-yellow/10 rounded-lg flex items-center justify-center">
<Eraser className="w-6 h-6 text-forge-yellow" />
</div>
<div>
<h1 className="text-2xl font-bold text-white">Background Remover</h1>
<p className="text-gray-500">Remove backgrounds instantly with AI precision</p>
</div>
</div>
<div className="grid grid-cols-1 lg:grid-cols-4 gap-8">
{/* Controls - Left Column */}
<div className="lg:col-span-1 space-y-6">
<div>
<label className="block text-sm font-medium text-gray-300 mb-2">
Upload Images
</label>
<FileUpload
onUploadMultiple={handleFileUpload}
accept={{ 'image/*': ['.png', '.jpg', '.jpeg', '.webp'] }}
label="Upload images (Multiple allowed)"
multiple={true}
/>
</div>
<div>
<label className="block text-sm font-medium text-gray-300 mb-2">
Output Format
</label>
<select
value={outputFormat}
onChange={(e: React.ChangeEvent<HTMLSelectElement>) => setOutputFormat(e.target.value)}
className="select-field"
>
{outputFormats.map((format) => (
<option key={format.value} value={format.value}>
{format.label}
</option>
))}
</select>
</div>
<div className="flex items-center gap-3">
<input
type="checkbox"
id="refineMask"
checked={refineMask}
onChange={(e: React.ChangeEvent<HTMLInputElement>) => setRefineMask(e.target.checked)}
className="w-4 h-4 rounded border-gray-600 bg-forge-dark text-forge-yellow focus:ring-forge-yellow"
/>
<label htmlFor="refineMask" className="text-gray-300 text-sm">
Refine edges (better quality, slower)
</label>
</div>
{queue.length > 0 && (
<div className="pt-4 border-t border-gray-800 space-y-4">
<button
onClick={handleProcessQueue}
disabled={processing || !queue.some(i => i.status === 'pending' || i.status === 'error')}
className="btn-primary w-full flex items-center justify-center gap-2"
>
{processing ? <RefreshCw className="w-4 h-4 animate-spin" /> : <Sparkles className="w-4 h-4" />}
{processing ? 'Processing Queue...' : 'Process All'}
</button>
<button
onClick={handleClearQueue}
disabled={processing}
className="btn-secondary w-full flex items-center justify-center gap-2"
>
<Trash2 className="w-4 h-4" /> Clear Queue
</button>
<div className="text-center text-xs text-gray-500">
{queue.filter(i => i.status === 'completed').length} / {queue.length} completed
</div>
</div>
)}
</div>
{/* Results / Queue - Right Column */}
<div className="lg:col-span-3 space-y-4">
<h2 className="text-lg font-semibold text-white">Queue</h2>
{queue.length === 0 ? (
<div className="bg-forge-dark rounded-xl border border-gray-800 p-12 flex flex-col items-center justify-center text-gray-500">
<Eraser className="w-12 h-12 mb-4 opacity-50" />
<p>Queue is empty. Upload images to start.</p>
</div>
) : (
<div className="space-y-4">
{queue.map(item => (
<div key={item.id} className="bg-forge-dark rounded-xl border border-gray-800 p-4 flex gap-4 items-start">
<div className="w-24 h-24 bg-black/40 rounded-lg flex-shrink-0 overflow-hidden relative border border-gray-700">
{/* Show source image if available (file or if we fetched asset details, but simpler to just show file or placeholder for now if from URL) */}
{item.file ? (
<img src={URL.createObjectURL(item.file)} alt="source" className="w-full h-full object-cover" />
) : item.assetId ? (
<img src={`/api/v1/assets/${item.assetId}/download`} alt="source" className="w-full h-full object-cover" />
) : (
<div className="w-full h-full flex items-center justify-center text-gray-600">No Img</div>
)}
</div>
<div className="flex-1 min-w-0">
<div className="flex justify-between items-start mb-2">
<h3 className="text-white font-medium truncate" title={item.originalFileName || item.assetId}>
{item.originalFileName || 'Asset ID: ' + item.assetId}
</h3>
<button onClick={() => removeItem(item.id)} className="text-gray-500 hover:text-red-400">
<Trash2 className="w-4 h-4" />
</button>
</div>
<div className="flex items-center gap-4">
{/* Status Badges */}
{item.status === 'pending' && <span className="text-xs font-medium px-2 py-1 bg-gray-800 text-gray-400 rounded">Pending</span>}
{item.status === 'uploading' && <span className="text-xs font-medium px-2 py-1 bg-blue-900/50 text-blue-400 rounded animate-pulse">Uploading...</span>}
{item.status === 'processing' && <span className="text-xs font-medium px-2 py-1 bg-yellow-900/50 text-forge-yellow rounded animate-pulse">Processing...</span>}
{item.status === 'error' && <span className="text-xs font-medium px-2 py-1 bg-red-900/50 text-red-400 rounded">Error: {item.error}</span>}
{item.status === 'completed' && <span className="text-xs font-medium px-2 py-1 bg-green-900/50 text-green-400 rounded">Completed</span>}
</div>
</div>
{/* Result Preview (if completed) */}
{item.status === 'completed' && item.resultAssetId && (
<div className="w-32 h-32 bg-[url('/grid.png')] bg-gray-800 rounded-lg flex-shrink-0 overflow-hidden relative border border-gray-700 bg-checker">
<img
src={`/api/v1/assets/${item.resultAssetId}/download`}
alt="result"
className="w-full h-full object-contain"
/>
<div className="absolute inset-0 bg-black/50 opacity-0 hover:opacity-100 transition-opacity flex items-center justify-center">
<button
onClick={() => handleDownload(item.resultAssetId!)}
className="p-2 bg-forge-yellow rounded-full hover:bg-yellow-400 text-black"
>
<Download className="w-4 h-4" />
</button>
</div>
</div>
)}
</div>
))}
</div>
)}
</div>
</div>
<style jsx global>{`
.bg-checker {
background-image: linear-gradient(45deg, #2a2a2a 25%, transparent 25%), linear-gradient(-45deg, #2a2a2a 25%, transparent 25%), linear-gradient(45deg, transparent 75%, #2a2a2a 75%), linear-gradient(-45deg, transparent 75%, #2a2a2a 75%);
background-size: 20px 20px;
background-position: 0 0, 0 10px, 10px -10px, -10px 0px;
}
`}</style>
</div>
);
}