540 lines
21 KiB
TypeScript
540 lines
21 KiB
TypeScript
"use client";
|
||
|
||
import React, { useCallback, useEffect, useMemo, useRef, useState } from "react";
|
||
import { Loader2 } from "lucide-react";
|
||
|
||
// ---------- Preview + batch helpers ----------
|
||
const PREVIEW_MAX = 2048; // max long-edge for on-screen previews
|
||
const BATCH_SIZES = [2048, 1536, 1280, 1024, 864, 720]; // adaptive preview long-edges
|
||
const COOLDOWN_MS = 150; // tiny cooldown between requests to ease VRAM
|
||
|
||
async function makePreview(blob: Blob, maxEdge = PREVIEW_MAX): Promise<Blob> {
|
||
const bitmap = await createImageBitmap(blob);
|
||
const { width, height } = bitmap;
|
||
const scale = Math.min(1, maxEdge / Math.max(width, height));
|
||
const outW = Math.max(1, Math.round(width * scale));
|
||
const outH = Math.max(1, Math.round(height * scale));
|
||
const canvas = document.createElement("canvas");
|
||
canvas.width = outW;
|
||
canvas.height = outH;
|
||
const ctx = canvas.getContext("2d")!;
|
||
ctx.drawImage(bitmap, 0, 0, outW, outH);
|
||
bitmap.close();
|
||
const outBlob = await new Promise<Blob>((resolve) =>
|
||
canvas.toBlob((b) => resolve(b!), "image/png")
|
||
);
|
||
return outBlob;
|
||
}
|
||
|
||
function revoke(url: string | null | undefined) {
|
||
if (!url) return;
|
||
try {
|
||
URL.revokeObjectURL(url);
|
||
} catch {}
|
||
}
|
||
|
||
// ---------- Methods ----------
|
||
type Canonical =
|
||
| "ormbg"
|
||
| "u2net"
|
||
| "basnet"
|
||
| "deeplab"
|
||
| "tracer"
|
||
| "u2net_human_seg"
|
||
| "isnet-general-use"
|
||
| "isnet-anime"
|
||
| "bria"
|
||
| "inspyrenet";
|
||
|
||
const METHODS: { key: Canonical; label: string }[] = [
|
||
{ key: "ormbg", label: "ORMBG" },
|
||
{ key: "u2net", label: "U2NET" },
|
||
{ key: "basnet", label: "BASNET" },
|
||
{ key: "deeplab", label: "DEEPLAB" },
|
||
{ key: "tracer", label: "TRACER-B7" },
|
||
{ key: "u2net_human_seg", label: "U2NET (Human)" },
|
||
{ key: "isnet-general-use", label: "ISNET (General)" },
|
||
{ key: "isnet-anime", label: "ISNET (Anime)" },
|
||
{ key: "bria", label: "BRIA RMBG1.4" },
|
||
{ key: "inspyrenet", label: "INSPYRENET" },
|
||
];
|
||
|
||
const DEFAULT_CONCURRENCY = 2;
|
||
|
||
type Status = "idle" | "pending" | "ok" | "error";
|
||
|
||
type ResultMap = {
|
||
[K in Canonical]?: {
|
||
fullBlob: Blob;
|
||
previewUrl: string;
|
||
bytes: number;
|
||
ms: number;
|
||
};
|
||
};
|
||
|
||
export default function BackgroundRemoverPage() {
|
||
// ---------- State ----------
|
||
const [file, setFile] = useState<File | null>(null);
|
||
const [sourceUrl, setSourceUrl] = useState<string | null>(null);
|
||
const [natural, setNatural] = useState<{ w: number; h: number } | null>(null);
|
||
|
||
const [status, setStatus] = useState<Record<Canonical, Status>>(
|
||
() => Object.fromEntries(METHODS.map((m) => [m.key, "idle"])) as Record<Canonical, Status>
|
||
);
|
||
const [results, setResults] = useState<ResultMap>({});
|
||
const resultsRef = useRef<ResultMap>({});
|
||
useEffect(() => {
|
||
resultsRef.current = results;
|
||
}, [results]);
|
||
|
||
const [active, setActive] = useState<Canonical | null>(null);
|
||
const [reveal, setReveal] = useState<number>(50);
|
||
const [gpuSafe, setGpuSafe] = useState(true);
|
||
|
||
const frameRef = useRef<HTMLDivElement | null>(null);
|
||
const draggingRef = useRef(false);
|
||
const batchBlobCache = useRef<Map<number, Blob>>(new Map());
|
||
|
||
useEffect(() => {
|
||
return () => {
|
||
revoke(sourceUrl);
|
||
Object.values(resultsRef.current).forEach((v) => revoke(v?.previewUrl));
|
||
};
|
||
}, [sourceUrl]);
|
||
|
||
// ---------- Styles ----------
|
||
const styles = (
|
||
<style>{`
|
||
html, body { width: 100%; overflow-x: hidden; }
|
||
:root { font-size: 17px; }
|
||
.checkerboard {
|
||
background-size: 24px 24px;
|
||
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-position: 0 0,0 12px,12px -12px,-12px 0;
|
||
}
|
||
.slider-handle { position: absolute; top: 0; bottom: 0; width: 0; left: calc(var(--reveal, 50) * 1%); }
|
||
.slider-handle::before { content: ""; position: absolute; top: 0; bottom: 0; width: 2px; left: -1px; background: rgba(255,255,255,0.85); }
|
||
.slider-thumb { position: absolute; top: 50%; transform: translate(-50%, -50%); width: 26px; height: 26px; border-radius: 9999px; background: rgba(24,24,27,0.9); border: 1px solid rgba(255,255,255,0.85); display: grid; place-items: center; cursor: ew-resize; }
|
||
/* Mobile: keep the page from panning left/right while using the slider */
|
||
.app-frame { touch-action: pan-y; overscroll-behavior-x: contain; }
|
||
`}</style>
|
||
);
|
||
|
||
// ---------- File pick ----------
|
||
const onPick = useCallback(
|
||
async (f: File | null) => {
|
||
revoke(sourceUrl);
|
||
Object.values(resultsRef.current).forEach((v) => revoke(v?.previewUrl));
|
||
batchBlobCache.current.clear();
|
||
|
||
setFile(f);
|
||
setResults({});
|
||
setActive(null);
|
||
setReveal(50);
|
||
setStatus(Object.fromEntries(METHODS.map((m) => [m.key, "idle"])) as any);
|
||
|
||
if (!f) {
|
||
setSourceUrl(null);
|
||
setNatural(null);
|
||
return;
|
||
}
|
||
|
||
try {
|
||
const bmp = await createImageBitmap(f);
|
||
setNatural({ w: bmp.width, h: bmp.height });
|
||
bmp.close();
|
||
} catch {}
|
||
|
||
const previewBlob = await makePreview(f, PREVIEW_MAX);
|
||
const previewUrl = URL.createObjectURL(previewBlob);
|
||
setSourceUrl(previewUrl);
|
||
},
|
||
[sourceUrl]
|
||
);
|
||
|
||
// Get or create a cached resized blob for batch preview runs
|
||
const getBatchBlob = useCallback(
|
||
async (longEdge: number): Promise<Blob> => {
|
||
const cache = batchBlobCache.current;
|
||
if (cache.has(longEdge)) return cache.get(longEdge)!;
|
||
if (!file) throw new Error("No file selected");
|
||
const b = await makePreview(file, longEdge);
|
||
cache.set(longEdge, b);
|
||
return b;
|
||
},
|
||
[file]
|
||
);
|
||
|
||
// ---------- Batch run (adaptive) ----------
|
||
const startAll = useCallback(async () => {
|
||
if (!file) return;
|
||
|
||
setResults({});
|
||
setStatus((prev) => {
|
||
const next = { ...prev };
|
||
METHODS.forEach((m) => (next[m.key] = "pending"));
|
||
return next;
|
||
});
|
||
|
||
const runOne = async (key: Canonical) => {
|
||
// When GPU-safe is on, try progressively smaller long-edge previews.
|
||
const sizes = gpuSafe ? BATCH_SIZES : [Math.max(natural?.w || 0, natural?.h || 0) || 4096];
|
||
|
||
let lastErr: string | null = null;
|
||
const t0 = performance.now();
|
||
|
||
for (const size of sizes) {
|
||
try {
|
||
const blobToSend = gpuSafe ? await getBatchBlob(size) : file!;
|
||
const fd = new FormData();
|
||
fd.append("file", blobToSend);
|
||
fd.append("method", key);
|
||
|
||
// CHANGED: call local proxy instead of a hardcoded service
|
||
const res = await fetch("/api/bgbye/process", { method: "POST", body: fd });
|
||
if (!res.ok) {
|
||
const txt = await res.text().catch(() => "");
|
||
const retryable = /out of memory|onnxruntime|cuda|allocate|500/i.test(txt);
|
||
if (gpuSafe && retryable) {
|
||
lastErr = txt || `HTTP ${res.status}`;
|
||
continue; // try next smaller size
|
||
}
|
||
throw new Error(txt || `HTTP ${res.status}`);
|
||
}
|
||
|
||
const outBlob = await res.blob();
|
||
const ms = performance.now() - t0;
|
||
const previewBlob = await makePreview(outBlob);
|
||
const previewUrl = URL.createObjectURL(previewBlob);
|
||
|
||
setResults((r) => ({
|
||
...r,
|
||
[key]: { fullBlob: outBlob, previewUrl, bytes: outBlob.size, ms },
|
||
}));
|
||
setStatus((s) => ({ ...s, [key]: "ok" }));
|
||
await new Promise((r) => setTimeout(r, COOLDOWN_MS)); // tiny cooldown
|
||
return;
|
||
} catch (e: any) {
|
||
lastErr = e?.message || String(e);
|
||
if (!gpuSafe) break; // not in adaptive mode, bail after first failure
|
||
}
|
||
}
|
||
|
||
setStatus((s) => ({ ...s, [key]: "error" }));
|
||
// (Optional) console.debug("Last error for", key, lastErr);
|
||
};
|
||
|
||
const concurrency = gpuSafe ? 1 : DEFAULT_CONCURRENCY;
|
||
const queue = [...METHODS.map((m) => m.key)];
|
||
let inFlight: Promise<void>[] = [];
|
||
|
||
const launch = () => {
|
||
while (inFlight.length < concurrency && queue.length) {
|
||
const key = queue.shift()!;
|
||
const p = runOne(key).finally(() => {
|
||
inFlight = inFlight.filter((q) => q !== p);
|
||
});
|
||
inFlight.push(p);
|
||
}
|
||
};
|
||
|
||
launch();
|
||
while (inFlight.length) {
|
||
await Promise.race(inFlight);
|
||
launch();
|
||
}
|
||
|
||
setActive((prev) => {
|
||
if (prev) return prev;
|
||
for (const m of METHODS) if ((resultsRef.current as any)[m.key]) return m.key;
|
||
return METHODS[0]?.key ?? null;
|
||
});
|
||
}, [file, gpuSafe, natural, getBatchBlob]);
|
||
|
||
// ---------- Upload & slider handlers ----------
|
||
const onDrop = (e: React.DragEvent) => {
|
||
e.preventDefault();
|
||
const f = e.dataTransfer.files?.[0];
|
||
if (f) onPick(f);
|
||
};
|
||
const onSelect = (e: React.ChangeEvent<HTMLInputElement>) => {
|
||
const f = e.target.files?.[0] ?? null;
|
||
onPick(f);
|
||
};
|
||
|
||
const aspect = useMemo(() => (!natural ? 16 / 9 : natural.w / natural.h || 16 / 9), [natural]);
|
||
|
||
const updateByClientX = useCallback((clientX: number) => {
|
||
const el = frameRef.current;
|
||
if (!el) return;
|
||
const rect = el.getBoundingClientRect();
|
||
const pct = ((clientX - rect.left) / rect.width) * 100;
|
||
setReveal(Math.min(100, Math.max(0, pct)));
|
||
}, []);
|
||
const onMouseDown = (e: React.MouseEvent) => {
|
||
draggingRef.current = true;
|
||
updateByClientX(e.clientX);
|
||
};
|
||
const onMouseMove = (e: React.MouseEvent) => {
|
||
if (!draggingRef.current) return;
|
||
updateByClientX(e.clientX);
|
||
};
|
||
const onMouseUp = () => (draggingRef.current = false);
|
||
const onTouchStart = (e: React.TouchEvent) => {
|
||
draggingRef.current = true;
|
||
updateByClientX(e.touches[0].clientX);
|
||
};
|
||
const onTouchMove = (e: React.TouchEvent) => {
|
||
if (!draggingRef.current) return;
|
||
e.preventDefault(); // keep page from horizontal panning while sliding
|
||
updateByClientX(e.touches[0].clientX);
|
||
};
|
||
const onTouchEnd = () => (draggingRef.current = false);
|
||
|
||
// ---------- Active result & actions ----------
|
||
const activeResult = active ? results[active] : undefined;
|
||
const canDownload = Boolean(active && activeResult?.fullBlob);
|
||
|
||
const download = () => {
|
||
if (!active || !activeResult) return;
|
||
const a = document.createElement("a");
|
||
const base = file?.name?.replace(/\.[^.]+$/, "") || "image";
|
||
const fullUrl = URL.createObjectURL(activeResult.fullBlob);
|
||
a.href = fullUrl;
|
||
setTimeout(() => revoke(fullUrl), 5000);
|
||
a.download = `${base}_${active}.png`;
|
||
document.body.appendChild(a);
|
||
a.click();
|
||
a.remove();
|
||
};
|
||
|
||
// Re-run the selected method on the ORIGINAL file for full-resolution output
|
||
const renderFullRes = useCallback(async () => {
|
||
if (!file || !active) return;
|
||
setStatus((s) => ({ ...s, [active]: "pending" }));
|
||
const t0 = performance.now();
|
||
try {
|
||
const fd = new FormData();
|
||
fd.append("file", file);
|
||
fd.append("method", active);
|
||
// CHANGED: call local proxy instead of a hardcoded service
|
||
const res = await fetch("/api/bgbye/process", { method: "POST", body: fd });
|
||
if (!res.ok) throw new Error(await res.text());
|
||
const outBlob = await res.blob();
|
||
const ms = performance.now() - t0;
|
||
const previewBlob = await makePreview(outBlob);
|
||
const previewUrl = URL.createObjectURL(previewBlob);
|
||
const prev = resultsRef.current[active];
|
||
if (prev) revoke(prev.previewUrl);
|
||
setResults((r) => ({ ...r, [active]: { fullBlob: outBlob, previewUrl, bytes: outBlob.size, ms } }));
|
||
setStatus((s) => ({ ...s, [active]: "ok" }));
|
||
} catch {
|
||
setStatus((s) => ({ ...s, [active]: "error" }));
|
||
}
|
||
}, [file, active]);
|
||
|
||
const doneCount = useMemo(() => METHODS.filter((m) => status[m.key] === "ok").length, [status]);
|
||
const pendingCount = useMemo(() => METHODS.filter((m) => status[m.key] === "pending").length, [status]);
|
||
|
||
function StatusDot({ s }: { s: Status }) {
|
||
const cls =
|
||
s === "ok"
|
||
? "bg-emerald-500"
|
||
: s === "pending"
|
||
? "bg-amber-400 animate-pulse"
|
||
: s === "error"
|
||
? "bg-rose-500"
|
||
: "bg-zinc-600";
|
||
return <span className={`inline-block w-2 h-2 rounded-full ${cls}`} />;
|
||
}
|
||
|
||
// ---------- Render ----------
|
||
return (
|
||
<div className="p-6 text-zinc-100 overflow-x-hidden">
|
||
{styles}
|
||
|
||
<div className="mx-auto w-full max-w-[1200px] px-4">
|
||
{/* Header row: title left, back button right */}
|
||
<div className="mb-4 flex items-center justify-between gap-3">
|
||
<h1 className="text-2xl font-semibold">Background Remover</h1>
|
||
<a
|
||
href="https://makearmy.io"
|
||
className="px-3 py-1 rounded-md border border-zinc-700 hover:bg-zinc-800/60 text-sm"
|
||
>
|
||
Back to main
|
||
</a>
|
||
</div>
|
||
|
||
{/* Source filename */}
|
||
<div className="text-zinc-400 mb-3">
|
||
<span className="text-zinc-300">Source:</span>{" "}
|
||
{file?.name ?? <span className="italic">— none —</span>}
|
||
</div>
|
||
|
||
{/* Preview frame (border removed per your note) */}
|
||
<div
|
||
ref={frameRef}
|
||
className="app-frame checkerboard relative w-full rounded-2xl shadow-inner"
|
||
style={{ aspectRatio: `${aspect}`, maxWidth: "1200px", maxHeight: "80vh", marginInline: "auto" }}
|
||
onDragOver={(e) => e.preventDefault()}
|
||
onDrop={onDrop}
|
||
onMouseDown={onMouseDown}
|
||
onMouseMove={onMouseMove}
|
||
onMouseLeave={onMouseUp}
|
||
onMouseUp={onMouseUp}
|
||
onTouchStart={onTouchStart}
|
||
onTouchMove={onTouchMove}
|
||
onTouchEnd={onTouchEnd}
|
||
>
|
||
{/* Drop hint */}
|
||
{!sourceUrl && (
|
||
<label className="absolute inset-0 grid place-items-center cursor-pointer">
|
||
<input type="file" accept="image/*" className="hidden" onChange={onSelect} />
|
||
<div className="text-zinc-400 border-2 border-dashed border-zinc-600/70 rounded-xl px-6 py-10">
|
||
<div className="text-center">
|
||
<div className="mb-1">Drop an image here</div>
|
||
<div className="text-zinc-500">or click to select a file</div>
|
||
</div>
|
||
</div>
|
||
</label>
|
||
)}
|
||
|
||
{/* Before/After */}
|
||
{sourceUrl && (
|
||
<>
|
||
{/* LEFT (BEFORE) */}
|
||
<img
|
||
src={sourceUrl}
|
||
loading="lazy"
|
||
decoding="async"
|
||
className="absolute inset-0 w-full h-full object-contain select-none"
|
||
alt="Source"
|
||
style={{ clipPath: `inset(0 0 0 ${reveal}%)` }}
|
||
draggable={false}
|
||
/>
|
||
|
||
{/* RIGHT (AFTER) */}
|
||
{activeResult ? (
|
||
<img
|
||
src={activeResult.previewUrl}
|
||
loading="lazy"
|
||
decoding="async"
|
||
className="absolute inset-0 w-full h-full object-contain select-none pointer-events-none"
|
||
alt="Result"
|
||
style={{ clipPath: `inset(0 ${100 - reveal}% 0 0)` }}
|
||
draggable={false}
|
||
/>
|
||
) : status[active as Canonical] === "pending" ? (
|
||
<div className="absolute inset-0 grid place-items-center">
|
||
<Loader2 className="animate-spin" />
|
||
</div>
|
||
) : null}
|
||
|
||
{/* Divider & Thumb */}
|
||
<div
|
||
className="slider-handle"
|
||
style={{ "--reveal": `${Math.min(100, Math.max(0, reveal))}` } as React.CSSProperties}
|
||
>
|
||
<div className="slider-thumb">
|
||
<div className="w-1.5 h-4 bg-white/80 rounded" />
|
||
</div>
|
||
</div>
|
||
</>
|
||
)}
|
||
</div>
|
||
|
||
{/* Method options – EVEN GRID under the preview */}
|
||
<div className="mt-4 grid grid-cols-2 sm:grid-cols-3 md:grid-cols-4 lg:grid-cols-5 gap-2">
|
||
{METHODS.map(({ key, label }) => (
|
||
<button
|
||
key={key}
|
||
className={`w-full justify-center px-3 py-2 rounded-md border flex items-center gap-2 ${
|
||
active === key ? "border-blue-400 bg-blue-500/20" : "border-zinc-700 hover:bg-zinc-800/60"
|
||
}`}
|
||
onClick={() => setActive(key)}
|
||
disabled={!file}
|
||
title={!file ? "Select a file first" : label}
|
||
>
|
||
<StatusDot s={status[key]} />
|
||
<span className="truncate">{label}</span>
|
||
</button>
|
||
))}
|
||
</div>
|
||
|
||
{/* Actions & global status */}
|
||
<div className="mt-4 flex items-center gap-3 flex-wrap">
|
||
<button
|
||
onClick={startAll}
|
||
className="px-3 py-1 rounded-md border border-zinc-700 hover:bg-zinc-800/60 flex items-center gap-2 order-0"
|
||
disabled={!file || pendingCount > 0}
|
||
title={!file ? "Select a file first" : pendingCount > 0 ? "Running…" : "Run all methods"}
|
||
>
|
||
{pendingCount > 0 && <Loader2 className="animate-spin w-4 h-4" />}{" "}
|
||
{pendingCount > 0 ? `Running… ${doneCount}/${METHODS.length}` : "Run all methods"}
|
||
</button>
|
||
|
||
{/* GPU-safe toggle */}
|
||
<label className="flex items-center gap-2 text-sm text-zinc-300 cursor-pointer select-none order-1">
|
||
<input type="checkbox" checked={gpuSafe} onChange={(e) => setGpuSafe(e.target.checked)} /> GPU-safe mode
|
||
</label>
|
||
|
||
<div className="text-zinc-400 text-sm order-2">
|
||
{file ? (
|
||
pendingCount > 0 ? (
|
||
<span>Processing… {doneCount}/{METHODS.length} finished</span>
|
||
) : doneCount > 0 ? (
|
||
<span>Done: {doneCount} methods succeeded</span>
|
||
) : (
|
||
<span>Ready. Click Run all methods</span>
|
||
)
|
||
) : (
|
||
<span>Drop an image to begin</span>
|
||
)}
|
||
</div>
|
||
|
||
{/* Right-side controls */}
|
||
<div className="sm:ml-auto flex items-center gap-3 w-full sm:w-auto order-3">
|
||
<input
|
||
type="range"
|
||
min={0}
|
||
max={100}
|
||
value={reveal}
|
||
onChange={(e) => setReveal(parseInt(e.target.value, 10))}
|
||
className="w-full sm:w-56"
|
||
title="Slide to compare before/after"
|
||
/>
|
||
|
||
<button
|
||
onClick={renderFullRes}
|
||
disabled={!file || !active}
|
||
className={`px-3 py-1 rounded-md border ${
|
||
file && active
|
||
? "border-sky-600 bg-sky-600/20 hover:bg-sky-600/30"
|
||
: "border-zinc-700 text-zinc-400 cursor-not-allowed"
|
||
}`}
|
||
title={!file ? "Select a file first" : !active ? "Choose a method" : "Render selected method at full resolution"}
|
||
>
|
||
Full-res render
|
||
</button>
|
||
|
||
<button
|
||
onClick={download}
|
||
disabled={!canDownload}
|
||
className={`px-3 py-1 rounded-md border ${
|
||
canDownload
|
||
? "border-emerald-600 bg-emerald-600/20 hover:bg-emerald-600/30"
|
||
: "border-zinc-700 text-zinc-400 cursor-not-allowed"
|
||
}`}
|
||
>
|
||
Download
|
||
</button>
|
||
</div>
|
||
</div>
|
||
</div>
|
||
</div>
|
||
);
|
||
}
|