diff --git a/PanTS-Demo/src/helpers/pendingUploads.ts b/PanTS-Demo/src/helpers/pendingUploads.ts new file mode 100644 index 0000000..1b61600 --- /dev/null +++ b/PanTS-Demo/src/helpers/pendingUploads.ts @@ -0,0 +1,106 @@ +// Resumable-upload store backed by IndexedDB. +// +// A chunked upload can be interrupted by a tab close/reload. The backend keeps +// every received chunk on disk (`/tmp/uploads//chunk-N`) and a +// re-sent chunk just overwrites, so the *server* side is already resumable - +// the only thing lost on reload is the browser's `File` object (in-memory refs +// don't survive a page unload). IndexedDB *does* persist Blobs across reloads, +// so we stash the picked file plus a chunk cursor here and resume on reopen. +// +// Everything is best-effort: if IndexedDB is unavailable or over quota, callers +// fall back to a normal (non-resumable) upload rather than failing the run. + +export type PendingUpload = { + sessionId: string; + file: Blob; // the picked File - File extends Blob, survives in IDB + filename: string; + model: string; + bdmapId: string; + totalChunks: number; + nextChunk: number; // first chunk not yet confirmed uploaded +}; + +const DB_NAME = "bodymaps-uploads"; +const STORE = "pending"; + +function openDb(): Promise { + return new Promise((resolve, reject) => { + const req = indexedDB.open(DB_NAME, 1); + req.onupgradeneeded = () => { + const db = req.result; + if (!db.objectStoreNames.contains(STORE)) { + db.createObjectStore(STORE, { keyPath: "sessionId" }); + } + }; + req.onsuccess = () => resolve(req.result); + req.onerror = () => reject(req.error); + }); +} + +function withStore( + mode: IDBTransactionMode, + fn: (store: IDBObjectStore) => IDBRequest | void, + onResult?: (req: IDBRequest) => T, +): Promise { + return openDb().then(db => + new Promise((resolve, reject) => { + const tx = db.transaction(STORE, mode); + const store = tx.objectStore(STORE); + const req = fn(store); + tx.oncomplete = () => { db.close(); resolve(req && onResult ? onResult(req) : undefined); }; + tx.onerror = () => { db.close(); reject(tx.error); }; + tx.onabort = () => { db.close(); reject(tx.error); }; + }) + ); +} + +export async function savePendingUpload(p: PendingUpload): Promise { + try { + await withStore("readwrite", store => store.put(p)); + return true; + } catch (e) { + console.warn("savePendingUpload failed - upload won't be resumable", e); + return false; + } +} + +export async function setPendingNextChunk(sessionId: string, nextChunk: number): Promise { + try { + const db = await openDb(); + await new Promise((resolve, reject) => { + const tx = db.transaction(STORE, "readwrite"); + const store = tx.objectStore(STORE); + const getReq = store.get(sessionId); + getReq.onsuccess = () => { + const rec = getReq.result as PendingUpload | undefined; + if (rec) { rec.nextChunk = nextChunk; store.put(rec); } + }; + tx.oncomplete = () => { db.close(); resolve(); }; + tx.onerror = () => { db.close(); reject(tx.error); }; + }); + } catch (e) { + console.warn("setPendingNextChunk failed", e); + } +} + +export async function deletePendingUpload(sessionId: string): Promise { + try { + await withStore("readwrite", store => store.delete(sessionId)); + } catch (e) { + console.warn("deletePendingUpload failed", e); + } +} + +export async function loadPendingUploads(): Promise { + try { + const result = await withStore( + "readonly", + store => store.getAll(), + req => (req.result as PendingUpload[]) || [], + ); + return (result as PendingUpload[]) || []; + } catch (e) { + console.warn("loadPendingUploads failed", e); + return []; + } +} diff --git a/PanTS-Demo/src/helpers/recentUploads.ts b/PanTS-Demo/src/helpers/recentUploads.ts index 37e4bb3..97a12bb 100644 --- a/PanTS-Demo/src/helpers/recentUploads.ts +++ b/PanTS-Demo/src/helpers/recentUploads.ts @@ -1,7 +1,7 @@ // Recent uploads persisted in the user's localStorage (mirrors JHU's recentIds // pattern). Extracted from UploadPage so the logic can be unit-tested. -export type RecentUploadStatus = "Processing" | "Completed" | "Failed"; +export type RecentUploadStatus = "Processing" | "Completed" | "Failed" | "Cancelled"; export type RecentUpload = { sessionId: string; @@ -39,6 +39,12 @@ export const addRecentUpload = (entry: RecentUpload): RecentUpload[] => { return trimmed; }; +export const removeRecentUpload = (sessionId: string): RecentUpload[] => { + const list = loadRecentUploads().filter((u) => u.sessionId !== sessionId); + persistRecentUploads(list); + return list; +}; + export const updateRecentUploadStatus = ( sessionId: string, status: RecentUploadStatus @@ -59,4 +65,10 @@ export const formatRelativeTime = (ts: number): string => { }; export const recentStatusColor = (status: RecentUploadStatus): string => - status === "Failed" ? "#ef4444" : status === "Processing" ? "#6a6a6a" : "#8f8f8f"; + status === "Failed" + ? "#ef4444" + : status === "Cancelled" + ? "#d97706" + : status === "Processing" + ? "#6a6a6a" + : "#8f8f8f"; diff --git a/PanTS-Demo/src/routes/UploadPage.css b/PanTS-Demo/src/routes/UploadPage.css index c37960f..c200874 100644 --- a/PanTS-Demo/src/routes/UploadPage.css +++ b/PanTS-Demo/src/routes/UploadPage.css @@ -450,6 +450,122 @@ color: #111111; } +/* Custom model dropdown */ +.model-dropdown { + position: relative; + width: 100%; +} +.model-dropdown-btn { + width: 100%; + padding: 10px 14px; + background: rgba(0,0,0,.04); + border: 1px solid rgba(0,0,0,.08); + border-radius: 8px; + color: #8f8f8f; + font-family: 'Space Grotesk', sans-serif; + font-size: 13px; + cursor: pointer; + outline: none; + transition: all 0.15s; + display: flex; + align-items: center; + justify-content: space-between; + text-align: left; +} +.model-dropdown-btn:hover { + border-color: rgba(0,0,0,.15); + background-color: rgba(0,0,0,.06); +} +.model-dropdown-btn.open { + border-color: rgba(0,0,0,.25); + box-shadow: 0 0 0 2px rgba(0,0,0,.06); +} +.model-dropdown-btn.has-value { + color: #111111; +} +.model-dropdown-chevron { + color: #6a6a6a; + flex-shrink: 0; + transition: transform 0.18s; +} +.model-dropdown-chevron.rotated { + transform: rotate(180deg); +} +.model-dropdown-menu { + position: absolute; + top: calc(100% + 4px); + left: 0; + right: 0; + background: #ffffff; + border: 1px solid rgba(0,0,0,.10); + border-radius: 10px; + box-shadow: 0 8px 24px rgba(0,0,0,.12); + z-index: 100; + overflow: hidden; +} +.model-dropdown-item { + display: flex; + align-items: center; + justify-content: space-between; + padding: 10px 14px; + cursor: pointer; + transition: background 0.1s; +} +.model-dropdown-item:hover { + background: rgba(0,0,0,.04); +} +.model-dropdown-item.selected { + background: rgba(0,45,114,.05); +} +.model-dropdown-item-content { + display: flex; + flex-direction: column; + gap: 1px; +} +.model-dropdown-item-name { + font-family: 'Space Grotesk', sans-serif; + font-size: 13px; + font-weight: 500; + color: #111111; +} +.model-dropdown-item.selected .model-dropdown-item-name { + color: #002D72; +} +.model-dropdown-item-desc { + font-family: 'Space Grotesk', sans-serif; + font-size: 11px; + color: #8f8f8f; +} +.model-dropdown-check { + color: #002D72; + flex-shrink: 0; +} +.model-dropdown-item-side { + display: flex; + align-items: center; + gap: 8px; + flex-shrink: 0; + margin-left: 10px; +} + +/* ── Cancel button on Active cards ── */ +.active-cancel-btn { + background: transparent; + border: 1px solid rgba(0, 0, 0, 0.1); + border-radius: 6px; + padding: 6px 12px; + color: #6a6a6a; + font-family: 'Space Grotesk', sans-serif; + font-size: 11px; + cursor: pointer; + transition: all 0.15s; +} +.active-cancel-btn:hover { + border-color: rgba(239, 68, 68, 0.4); + color: #ef4444; + background: rgba(239, 68, 68, 0.05); +} + /* Arrow connector */ .pipeline-arrow { display: flex; @@ -981,3 +1097,16 @@ min-width: 54px; text-align: right; } + +/* ── Active-upload spinner ── */ +@keyframes upload-spin { + to { transform: rotate(360deg); } +} +.upload-spinner { + width: 18px; + height: 18px; + border-radius: 50%; + border: 2px solid rgba(0, 45, 114, 0.12); + border-top-color: #002D72; + animation: upload-spin 0.75s linear infinite; +} diff --git a/PanTS-Demo/src/routes/UploadPage.tsx b/PanTS-Demo/src/routes/UploadPage.tsx index 7624dc1..1f1b528 100644 --- a/PanTS-Demo/src/routes/UploadPage.tsx +++ b/PanTS-Demo/src/routes/UploadPage.tsx @@ -1,131 +1,46 @@ import React, { lazy, Suspense, useCallback, useEffect, useRef, useState } from 'react'; + +const MODEL_OPTIONS: { id: string; label: string; desc: string }[] = [ + { id: "ePAI", label: "ePAI", desc: "For detailed pancreas and tumor analysis" }, + { id: "SuPreM", label: "SuPreM", desc: "For whole-body scans from lungs to legs" }, + { id: "MedFormer", label: "MedFormer", desc: "For reliable abdominal segmentation" }, + { id: "R-Super", label: "R-Super", desc: "For the highest tumor detection accuracy" }, + { id: "Atlas-Net", label: "Atlas-Net", desc: "For anatomically consistent results" }, +]; import { useNavigate } from 'react-router-dom'; import './UploadPage.css'; -import { API_BASE } from '../helpers/constants'; -import Header from '../components/Header'; +// Lazy so NiiVue isn't pulled into the upload bundle until a file is actually previewed. const CtPreview = lazy(() => import('../components/CtPreview/CtPreview')); -const DicomPreview = lazy(() => import('../components/CtPreview/DicomPreview')); - -const CHUNK_SIZE = 256 * 1024; -const NIFTI_EXTS = ['.nii', '.nii.gz']; -const DICOM_EXTS = ['.dcm', '.dicom']; - -const PIPELINE_INFO: Record = { - // Preprocessing - OpenVAE: 'A 3D variational autoencoder pretrained on large CT datasets. Reconstructs and denoises the input scan before segmentation to improve downstream model accuracy.', - // Models - ePAI: 'Pancreas AI — an nnU-Net based model specialized for pancreatic segmentation including the pancreas, ducts, and pancreatic lesions (PDAC, cysts, PNETs).', - SuPreM: 'Universal multi-organ segmentation model pretrained on 25 abdominal structures. Uses a UNet backbone with large-scale supervised pretraining across multiple datasets.', - MedFormer: 'Transformer-based segmentation model for 26 abdominal structures plus pancreatic lesions. Leverages attention mechanisms for long-range spatial context.', - 'R-Super': 'Report-supervised extension of MedFormer. Incorporates radiology report text during training to improve segmentation accuracy on underrepresented structures.', - 'Atlas-Net': 'Atlas-guided nnU-Net model that uses anatomical priors to improve robustness across varied CT acquisition protocols and patient populations.', - // Postprocessing - ShapeKit: 'CPU-only shape refinement toolkit. Uses anatomical shape priors and connected-component analysis to clean up segmentation boundaries and remove spurious predictions.', -}; - -type JobStatus = 'queued' | 'uploading' | 'processing' | 'completed' | 'failed'; - -type PendingItem = { - id: string; - displayName: string; - files: File[]; - isDicom: boolean; -}; - -type UploadJob = { - id: string; - displayName: string; - files: File[]; - isDicom: boolean; - sessionId: string | null; - status: JobStatus; - uploadProgress: number; - inferenceProgress: number; - model: string; - timestamp: number; - error?: string; -}; +import { API_BASE } from '../helpers/constants'; +import { + addRecentUpload, + formatRelativeTime, + loadRecentUploads, + recentStatusColor, + removeRecentUpload, + updateRecentUploadStatus, + type RecentUpload, +} from '../helpers/recentUploads'; +import Header from '../components/Header'; +import { looksLikeDicom, setLocalDicomFiles } from '../helpers/dicomLocal'; +import { + deletePendingUpload, + loadPendingUploads, + savePendingUpload, + setPendingNextChunk, + type PendingUpload, +} from '../helpers/pendingUploads'; const parseApiResponse = async (res: Response): Promise => { - const ct = res.headers.get('content-type') || ''; - if (ct.includes('application/json')) return res.json(); - const text = await res.text(); - throw new Error(`HTTP ${res.status}: ${text.slice(0, 200).replace(/\s+/g, ' ').trim()}`); -}; - -// Completed/failed results survive a tab reload: the queue is mirrored to localStorage -// (minus the File objects, which can't be serialized). Jobs that were still in flight on -// reload — uploading, queued, or running inference — are NOT re-attached or resumed; -// they're marked failed so reopening the page never auto-resumes any work. -const STORAGE_KEY = 'pants_upload_jobs_v1'; - -const loadPersistedJobs = (): UploadJob[] => { - try { - const raw = localStorage.getItem(STORAGE_KEY); - if (!raw) return []; - const saved = JSON.parse(raw) as Omit[]; - return saved.map(j => - j.status === 'uploading' || j.status === 'queued' || j.status === 'processing' - ? { ...j, files: [], status: 'failed' as JobStatus, error: 'Interrupted by page reload — please re-upload' } - : { ...j, files: [] } - ); - } catch { - return []; + const contentType = res.headers.get("content-type") || ""; + if (contentType.includes("application/json")) { + return res.json(); } -}; - -type PipelineOption = { value: string; label: string }; - -const PipelineSelect: React.FC<{ - value: string; - onChange: (v: string) => void; - options: PipelineOption[]; - placeholder?: string; -}> = ({ value, onChange, options, placeholder }) => { - const [open, setOpen] = useState(false); - const ref = useRef(null); - - useEffect(() => { - const handler = (e: MouseEvent) => { - if (ref.current && !ref.current.contains(e.target as Node)) setOpen(false); - }; - document.addEventListener('mousedown', handler); - return () => document.removeEventListener('mousedown', handler); - }, []); - - const selectedLabel = options.find(o => o.value === value)?.label; - - return ( -
- - {open && ( -
- {options.map(opt => ( -
{ onChange(opt.value); setOpen(false); }} - > - {opt.label} - {PIPELINE_INFO[opt.value] && ( - e.stopPropagation()}> - i -
{PIPELINE_INFO[opt.value]}
- - )} -
- ))} -
- )} -
+ const text = await res.text(); + const shortBody = text.slice(0, 200).replace(/\s+/g, " ").trim(); + throw new Error( + `Expected JSON but got ${contentType || "unknown content-type"} (HTTP ${res.status}). Body: ${shortBody}` ); }; @@ -133,93 +48,90 @@ const UploadPage: React.FC = () => { const navigate = useNavigate(); const fileInputRef = useRef(null); const dicomInputRef = useRef(null); - const jobQueueRef = useRef([]); - const isProcessingRef = useRef(false); - const cancelledRef = useRef(false); - - const [pendingItems, setPendingItems] = useState([]); - const [previewItemId, setPreviewItemId] = useState(null); - const [selectedModel, setSelectedModel] = useState(''); - const [selectedPreprocessing, setSelectedPreprocessing] = useState(''); - const [selectedPostprocessing, setSelectedPostprocessing] = useState(''); - const [isDragOver, setIsDragOver] = useState(false); - const [jobs, setJobs] = useState(loadPersistedJobs); - - /* ── File selection ── */ - const addFiles = (files: File[]) => { - const niftiFiles = files.filter(f => - NIFTI_EXTS.some(ext => f.name.toLowerCase().endsWith(ext)) - ); - const dicomFiles = files.filter(f => - DICOM_EXTS.some(ext => f.name.toLowerCase().endsWith(ext)) - ); - - const newItems: PendingItem[] = []; - - niftiFiles.forEach(f => newItems.push({ - id: crypto.randomUUID(), - displayName: f.name, - files: [f], - isDicom: false, - })); - - if (dicomFiles.length > 0) { - newItems.push({ - id: crypto.randomUUID(), - displayName: `DICOM Series (${dicomFiles.length} slices)`, - files: dicomFiles, - isDicom: true, - }); - } - - if (newItems.length === 0) { - alert('Please select .nii, .nii.gz, or .dcm files only.'); + // One poll timer per in-flight session so runs can proceed in parallel. + const pollTimersRef = useRef>>(new Map()); + // Whether the current foreground upload got stored in IndexedDB (resumable). + // If IDB was unavailable we fall back to warning before an unload instead. + const uploadResumableRef = useRef(false); + // AbortController per session so a mid-upload run can be cancelled cleanly. + const uploadAbortRef = useRef>(new Map()); + // Which session currently drives the foreground upload progress bar. + const foregroundUploadSidRef = useRef(null); + + // Local DICOM: stash the picked folder's files and open the viewer's /dicom + // route. Nothing is uploaded - the viewer reads the File objects directly. + const handleDicomFolderSelect = (e: React.ChangeEvent) => { + const files = Array.from(e.target.files ?? []); + e.target.value = ""; // allow re-picking the same folder later + const candidates = files.filter(looksLikeDicom); + if (!candidates.length) { + alert("No DICOM files (.dcm) found in the selected folder."); return; } - setPendingItems(prev => [...prev, ...newItems]); + setLocalDicomFiles(candidates); + navigate("/dicom"); }; + const [selectedFiles, setSelectedFiles] = useState([]); + const [message, setMessage] = useState(""); + const [serverPath, setServerPath] = useState(""); + const [sessionId, setSessionId] = useState(""); + const [bdmapId, setBdmapId] = useState(""); + const [uploadProgress, setUploadProgress] = useState(0); + const [isUploading, setIsUploading] = useState(false); + const [inferenceCompleted, setInferenceCompleted] = useState(false); + const [selectedModel, setSelectedModel] = useState<"ePAI" | "SuPreM" | "OpenVAE" | "MedFormer" | "R-Super" | "Atlas-Net" | "">(""); + const [modelDropOpen, setModelDropOpen] = useState(false); + const modelDropRef = useRef(null); + const [preDropOpen, setPreDropOpen] = useState(false); + const preDropRef = useRef(null); + const [preValue, setPreValue] = useState(""); + const [postDropOpen, setPostDropOpen] = useState(false); + const postDropRef = useRef(null); + const [postValue, setPostValue] = useState(""); + const [showAdvanced, setShowAdvanced] = useState(false); + const [isDragOver, setIsDragOver] = useState(false); + const [recentUploads, setRecentUploads] = useState(() => loadRecentUploads()); + // Sub-state of each Active card: "uploading" | "queued" | "running". + const [sessionPhases, setSessionPhases] = useState>({}); + + const setPhase = (sid: string, phase?: string) => + setSessionPhases(prev => { + if (phase === undefined) { + if (!(sid in prev)) return prev; + const { [sid]: _dropped, ...rest } = prev; + return rest; + } + return prev[sid] === phase ? prev : { ...prev, [sid]: phase }; + }); + + const allowedExtensions = [".nii", ".nii.gz"]; + + /* ── File handling ── */ const handleFileSelect = (e: React.ChangeEvent) => { if (!e.target.files) return; - addFiles(Array.from(e.target.files)); - e.target.value = ''; + const filteredFiles = Array.from(e.target.files).filter(file => + allowedExtensions.some(ext => file.name.toLowerCase().endsWith(ext)) + ); + if (filteredFiles.length === 0) { + alert("Please select .nii or .nii.gz files only"); + return; + } + setSelectedFiles(prev => [...prev, ...filteredFiles]); }; - const handleDrop = useCallback(async (e: React.DragEvent) => { + const handleDrop = useCallback((e: React.DragEvent) => { e.preventDefault(); setIsDragOver(false); - - const readEntry = (entry: FileSystemEntry): Promise => { - if (entry.isFile) { - return new Promise(res => (entry as FileSystemFileEntry).file(f => res([f]), () => res([]))); - } - if (entry.isDirectory) { - const reader = (entry as FileSystemDirectoryEntry).createReader(); - return new Promise(res => { - const all: FileSystemEntry[] = []; - const readBatch = () => - reader.readEntries(batch => { - if (!batch.length) { - Promise.all(all.map(readEntry)).then(groups => res(groups.flat())).catch(() => res([])); - } else { - all.push(...batch); - readBatch(); - } - }, () => res([])); - readBatch(); - }); - } - return Promise.resolve([]); - }; - - const items = Array.from(e.dataTransfer.items); - const entries = items.map(i => i.webkitGetAsEntry()).filter(Boolean) as FileSystemEntry[]; - if (entries.length > 0) { - const files = (await Promise.all(entries.map(readEntry))).flat(); - addFiles(files); - } else if (e.dataTransfer.files.length > 0) { - addFiles(Array.from(e.dataTransfer.files)); + if (!e.dataTransfer.files) return; + const filteredFiles = Array.from(e.dataTransfer.files).filter(file => + allowedExtensions.some(ext => file.name.toLowerCase().endsWith(ext)) + ); + if (filteredFiles.length === 0) { + alert("Please drop .nii or .nii.gz files only"); + return; } + setSelectedFiles(prev => [...prev, ...filteredFiles]); }, []); const handleDragOver = useCallback((e: React.DragEvent) => { @@ -227,246 +139,471 @@ const UploadPage: React.FC = () => { setIsDragOver(true); }, []); - const handleDragLeave = useCallback(() => setIsDragOver(false), []); + const handleDragLeave = useCallback(() => { + setIsDragOver(false); + }, []); - const removePendingItem = (id: string) => { - setPendingItems(prev => prev.filter(item => item.id !== id)); - if (previewItemId === id) setPreviewItemId(null); + const removeFile = (index: number) => { + setSelectedFiles(prev => prev.filter((_, i) => i !== index)); + }; + + /* ── Inference polling (one timer per session) ── */ + const stopPolling = (sid: string) => { + const timer = pollTimersRef.current.get(sid); + if (timer) { + clearInterval(timer); + pollTimersRef.current.delete(sid); + } }; - const togglePreview = (id: string) => - setPreviewItemId(prev => (prev === id ? null : id)); - - /* ── Job state helpers ── */ - const patchJob = (id: string, patch: Partial) => - setJobs(prev => prev.map(j => (j.id === id ? { ...j, ...patch } : j))); - - /* ── Polling ── */ - const pollUntilDone = (sessionId: string, jobId: string): Promise => - new Promise(resolve => { - const interval = setInterval(async () => { - try { - const res = await fetch(`${API_BASE}/api/inference-status/${sessionId}`); - const data = await parseApiResponse(res); - const status = (data.status || '').toLowerCase(); - if (status === 'completed') { - clearInterval(interval); - patchJob(jobId, { status: 'completed', inferenceProgress: 100 }); - resolve(); - } else if (status === 'failed') { - clearInterval(interval); - patchJob(jobId, { status: 'failed', error: data.error || 'Inference failed' }); - resolve(); - } else { - setJobs(prev => prev.map(j => - j.id === jobId - ? { ...j, inferenceProgress: Math.min(95, j.inferenceProgress + 7) } - : j - )); + const stopAllPolling = () => { + pollTimersRef.current.forEach(timer => clearInterval(timer)); + pollTimersRef.current.clear(); + }; + + const finishSession = (sid: string, model: string) => { + stopPolling(sid); + setPhase(sid); + const uploads = updateRecentUploadStatus(sid, "Completed"); + setRecentUploads(uploads); + setSessionId(sid); + setInferenceCompleted(true); + // Only auto-open the viewer when no other run is still in flight. + if (!uploads.some(u => u.status === "Processing")) { + setTimeout(() => { + navigate(model === "OpenVAE" ? `/reconstruction/${sid}` : `/session/${sid}`); + }, 600); + } + }; + + const startInferencePolling = (sid: string, model: string) => { + stopPolling(sid); + let notFoundCount = 0; + const timer = setInterval(async () => { + try { + const res = await fetch(`${API_BASE}/api/inference-status/${sid}`); + const data = await parseApiResponse(res); + const status = (data.status || "").toLowerCase(); + + // The server doesn't know this session: the upload never finished + // (tab closed mid-upload) or the backend restarted and lost its + // in-memory job table. A few consecutive hits = gone, not a blip. + if (status === "not_found") { + notFoundCount += 1; + if (notFoundCount >= 3) { + stopPolling(sid); + setPhase(sid); + setRecentUploads(updateRecentUploadStatus(sid, "Failed")); + setMessage("Session no longer exists on the server - marked as Failed."); } - } catch { /* keep polling */ } - }, 2500); - }); + return; + } + notFoundCount = 0; + + if (!res.ok) throw new Error(data.error || data.status || "Status check failed"); + + if (status === "completed") { + finishSession(sid, model); + } else if (status === "failed") { + stopPolling(sid); + setPhase(sid); + setRecentUploads(updateRecentUploadStatus(sid, "Failed")); + setMessage(`Inference failed${data.error ? `: ${data.error}` : ""}`); + } else if (status === "cancelled") { + // Cancelled elsewhere (another tab, or the backend) - reflect it. + stopPolling(sid); + setPhase(sid); + setRecentUploads(updateRecentUploadStatus(sid, "Cancelled")); + } else if (status === "queued" || status === "running") { + setPhase(sid, status); + } + } catch (err) { + // Network blip or proxy error while the backend restarts - the job + // may still be alive server-side, so keep polling. + console.error(err); + } + }, 2500); + pollTimersRef.current.set(sid, timer); + }; + + // Cancel one run, whatever phase it's in: aborts an in-flight upload or + // kills a queued/running server job. + const cancelRun = (upload: RecentUpload) => { + const sid = upload.sessionId; + stopPolling(sid); + setPhase(sid); + + const controller = uploadAbortRef.current.get(sid); + if (controller) controller.abort(); + deletePendingUpload(sid); + + // Fire-and-forget: if the job never reached the server (upload phase) + // this 404s, which is fine - the client side is already torn down. + fetch(`${API_BASE}/api/cancel-inference/${sid}`, { method: "POST" }).catch(() => {}); + + if (foregroundUploadSidRef.current === sid) { + foregroundUploadSidRef.current = null; + setIsUploading(false); + } + setRecentUploads(updateRecentUploadStatus(sid, "Cancelled")); + setMessage(`Cancelled ${upload.label}`); + }; - /* ── Persistence: mirror the queue to localStorage on every change ── */ useEffect(() => { + // Resume every in-flight run - there can be several in parallel. Uploads + // that were still mid-transfer live in IndexedDB and must be *resumed* + // (not polled - the server has no job for them yet); the rest are already + // inferencing server-side, so we reconnect their pollers. + let cancelled = false; + (async () => { + const processing = loadRecentUploads().filter(u => u.status === "Processing"); + const pending = await loadPendingUploads(); + if (cancelled) return; + const pendingById = new Map(pending.map(p => [p.sessionId, p])); + + for (const u of processing) { + if (pendingById.has(u.sessionId)) { + runUpload(pendingById.get(u.sessionId)!, false); // resume the upload + } else { + startInferencePolling(u.sessionId, u.model); // resume polling + } + } + + // Clean up IndexedDB entries whose card no longer exists (deleted or + // trimmed off the 8-entry list) so the store can't leak. + const known = new Set(loadRecentUploads().map(u => u.sessionId)); + pending.filter(p => !known.has(p.sessionId)).forEach(p => deletePendingUpload(p.sessionId)); + + if (processing.length > 0) { + setSessionId(processing[0].sessionId); + setMessage(`Reconnected · ${processing.length} run${processing.length === 1 ? "" : "s"} in progress`); + } + })(); + return () => { cancelled = true; stopAllPolling(); }; + }, []); + + // Only warn before an unload if the current upload could NOT be stored in + // IndexedDB (quota/private-mode) - otherwise an interrupted upload resumes + // automatically on reopen, so no scary dialog is needed. + useEffect(() => { + if (!isUploading || uploadResumableRef.current) return; + const handler = (e: BeforeUnloadEvent) => { + e.preventDefault(); + e.returnValue = ""; + }; + window.addEventListener("beforeunload", handler); + return () => window.removeEventListener("beforeunload", handler); + }, [isUploading]); + + useEffect(() => { + if (!modelDropOpen) return; + const handler = (e: MouseEvent) => { + if (modelDropRef.current && !modelDropRef.current.contains(e.target as Node)) + setModelDropOpen(false); + }; + document.addEventListener("mousedown", handler); + return () => document.removeEventListener("mousedown", handler); + }, [modelDropOpen]); + + useEffect(() => { + if (!preDropOpen) return; + const handler = (e: MouseEvent) => { + if (preDropRef.current && !preDropRef.current.contains(e.target as Node)) + setPreDropOpen(false); + }; + document.addEventListener("mousedown", handler); + return () => document.removeEventListener("mousedown", handler); + }, [preDropOpen]); + + useEffect(() => { + if (!postDropOpen) return; + const handler = (e: MouseEvent) => { + if (postDropRef.current && !postDropRef.current.contains(e.target as Node)) + setPostDropOpen(false); + }; + document.addEventListener("mousedown", handler); + return () => document.removeEventListener("mousedown", handler); + }, [postDropOpen]); + + /* ── Upload (chunked) ── */ + const CHUNK_SIZE = 256 * 1024; + + // Uploads the file described by `p` from p.nextChunk onward, finalizes, then + // starts inference. Resumable: the file lives in IndexedDB and the chunk + // cursor is persisted, so a reload can call this again to pick up where it + // left off. `foreground` = the run the user just clicked (drives the progress + // bar); resumed background runs show only their Active-section spinner. + const runUpload = async (p: PendingUpload, foreground: boolean) => { + const { sessionId: sid, file, filename, model, bdmapId: bid, totalChunks } = p; + const controller = new AbortController(); + uploadAbortRef.current.set(sid, controller); + setPhase(sid, "uploading"); try { - const persistable = jobs.map(j => ({ - id: j.id, - displayName: j.displayName, - isDicom: j.isDicom, - sessionId: j.sessionId, - status: j.status, - uploadProgress: j.uploadProgress, - inferenceProgress: j.inferenceProgress, - model: j.model, - timestamp: j.timestamp, - error: j.error, - })); - localStorage.setItem(STORAGE_KEY, JSON.stringify(persistable)); - } catch { /* storage unavailable — non-fatal */ } - }, [jobs]); - - /* ── Process NIfTI job ── */ - const processNiftiJob = async (job: UploadJob): Promise => { - const sid = crypto.randomUUID(); - patchJob(job.id, { sessionId: sid, status: 'uploading', uploadProgress: 0 }); - const file = job.files[0]; - const totalChunks = Math.ceil(file.size / CHUNK_SIZE); - try { - for (let i = 0; i < totalChunks; i++) { + if (foreground) { + foregroundUploadSidRef.current = sid; + setIsUploading(true); + setUploadProgress(Math.round((p.nextChunk / totalChunks) * 100)); + setMessage(`Uploading ${filename}...`); + } + + for (let i = p.nextChunk; i < totalChunks; i++) { const chunk = file.slice(i * CHUNK_SIZE, Math.min((i + 1) * CHUNK_SIZE, file.size)); - const fd = new FormData(); - fd.append('session_id', sid); - fd.append('chunk_index', i.toString()); - fd.append('total_chunks', totalChunks.toString()); - fd.append('file', chunk); - const res = await fetch(`${API_BASE}/api/upload-inference-chunk`, { method: 'POST', body: fd }); + const formData = new FormData(); + formData.append("session_id", sid); + formData.append("chunk_index", i.toString()); + formData.append("total_chunks", totalChunks.toString()); + formData.append("file", chunk); + + const res = await fetch(`${API_BASE}/api/upload-inference-chunk`, { + method: "POST", body: formData, signal: controller.signal, + }); + if (res.status === 413) throw new Error("Upload chunk too large for server/proxy limit (HTTP 413)."); const data = await parseApiResponse(res); - if (!res.ok) throw new Error(data.error || 'Chunk upload failed'); - patchJob(job.id, { uploadProgress: Math.round(((i + 1) / totalChunks) * 100) }); + if (!res.ok) throw new Error(data.error || "Chunk upload failed"); + + // Persist the cursor every so often (not every chunk - that would be a + // lot of IDB writes). On resume we re-send at most a few already-stored + // chunks, which the backend just overwrites. Harmless. + if (i % 16 === 0) await setPendingNextChunk(sid, i + 1); + if (foreground) setUploadProgress(Math.round(((i + 1) / totalChunks) * 100)); } + if (foreground) setMessage("Finalizing upload..."); const finalizeRes = await fetch(`${API_BASE}/api/finalize-upload`, { - method: 'POST', + method: "POST", + signal: controller.signal, body: new URLSearchParams({ session_id: sid, total_chunks: totalChunks.toString(), - output_filename: file.name, + output_filename: filename, + ...(bid ? { bdmap_id: bid } : {}), }), }); const finalizeData = await parseApiResponse(finalizeRes); - if (!finalizeRes.ok) throw new Error(finalizeData.error || 'Finalize failed'); - - patchJob(job.id, { status: 'processing', inferenceProgress: 5 }); + if (!finalizeRes.ok) throw new Error(finalizeData.error); + const uploadedName = finalizeData.uploaded_filename || filename; + + // File is fully on the server now - drop the IDB copy before we kick off + // inference so a later reload resumes by polling, not re-uploading. + await deletePendingUpload(sid); + if (foreground) { + foregroundUploadSidRef.current = null; + setUploadProgress(100); + setIsUploading(false); + } + if (foreground) setMessage(`Starting ${model} inference...`); const inferFd = new FormData(); - inferFd.append('session_id', sid); - inferFd.append('model_name', job.model); - inferFd.append('uploaded_filename', finalizeData.uploaded_filename || file.name); - const inferRes = await fetch(`${API_BASE}/api/run-epai-inference`, { method: 'POST', body: inferFd }); - const inferData = await parseApiResponse(inferRes); - if (!inferRes.ok) throw new Error(inferData.error || 'Failed to start inference'); - - const actualSid = inferData.session_id || sid; - if (actualSid !== sid) patchJob(job.id, { sessionId: actualSid }); - await pollUntilDone(actualSid, job.id); + inferFd.append("session_id", sid); + inferFd.append("model_name", model); + inferFd.append("uploaded_filename", uploadedName); + const res = await fetch(`${API_BASE}/api/run-epai-inference`, { + method: "POST", body: inferFd, signal: controller.signal, + }); + const data = await parseApiResponse(res); + if (!res.ok) throw new Error(data.error || "Failed to start inference"); + + setSessionId(sid); + setPhase(sid, "queued"); // server queues for the GPU; poll refines this + if (foreground) setMessage(`${model} inference started. Session: ${sid}`); + startInferencePolling(sid, model); } catch (err) { - patchJob(job.id, { status: 'failed', error: (err as Error).message }); + // A user cancel aborts our fetches - cancelRun already did the cleanup + // and set the card to Cancelled, so don't overwrite that with Failed. + if (controller.signal.aborted) return; + console.error(err); + setPhase(sid); + if (foreground) { + foregroundUploadSidRef.current = null; + setIsUploading(false); + } + await deletePendingUpload(sid); + setRecentUploads(updateRecentUploadStatus(sid, "Failed")); + if (foreground) setMessage("Failed: " + (err as Error).message); + } finally { + if (uploadAbortRef.current.get(sid) === controller) { + uploadAbortRef.current.delete(sid); + } } }; - /* ── Process DICOM job ── */ - const processDicomJob = async (job: UploadJob): Promise => { + /* ── Run inference ── */ + const handleRunEpaiInference = async () => { + const file = selectedFiles[0] ?? null; + const path = serverPath.trim(); + if (!file && !path) { + alert("Provide a server file path or upload/select a file first."); + return; + } + + const model = selectedModel; const sid = crypto.randomUUID(); - patchJob(job.id, { sessionId: sid, status: 'uploading', uploadProgress: 0 }); - try { - const total = job.files.length; - for (let i = 0; i < total; i++) { + const label = bdmapId.trim() || (path ? path.split("/").pop() : file?.name) || sid; + + setInferenceCompleted(false); + setRecentUploads( + addRecentUpload({ + sessionId: sid, + label, + model, + status: "Processing", + timestamp: Date.now(), + isReconstruction: model === "OpenVAE", + }) + ); + + // Server-path run: nothing to upload, kick off inference directly. A server + // path always wins over a selected file, matching the prior behavior. + if (path) { + try { + setSessionId(sid); + setMessage(`Starting ${model} inference...`); const fd = new FormData(); - fd.append('session_id', sid); - fd.append('file', job.files[i]); - const res = await fetch(`${API_BASE}/api/upload-dicom-slice`, { method: 'POST', body: fd }); + fd.append("session_id", sid); + fd.append("model_name", model); + fd.append("INPUT_SERVER_PATH", path); + const res = await fetch(`${API_BASE}/api/run-epai-inference`, { method: "POST", body: fd }); const data = await parseApiResponse(res); - if (!res.ok) throw new Error(data.error || 'DICOM slice upload failed'); - patchJob(job.id, { uploadProgress: Math.round(((i + 1) / total) * 100) }); + if (!res.ok) throw new Error(data.error || "Failed to start inference"); + setMessage(`${model} inference started. Session: ${sid}`); + setPhase(sid, "queued"); + startInferencePolling(sid, model); + } catch (err) { + console.error(err); + setRecentUploads(updateRecentUploadStatus(sid, "Failed")); + setMessage("Failed: " + (err as Error).message); } + return; + } - const finalizeRes = await fetch(`${API_BASE}/api/finalize-dicom`, { - method: 'POST', - body: new URLSearchParams({ session_id: sid }), - }); - const finalizeData = await parseApiResponse(finalizeRes); - if (!finalizeRes.ok) throw new Error(finalizeData.error || 'DICOM conversion failed'); - - patchJob(job.id, { status: 'processing', inferenceProgress: 5 }); + // File upload: consume it so the next file can be queued, stash it in + // IndexedDB so an interrupted upload can resume, then run. + setSelectedFiles(prev => prev.slice(1)); + const pending: PendingUpload = { + sessionId: sid, + file: file!, + filename: file!.name, + model, + bdmapId: bdmapId.trim(), + totalChunks: Math.ceil(file!.size / CHUNK_SIZE), + nextChunk: 0, + }; + uploadResumableRef.current = await savePendingUpload(pending); + runUpload(pending, true); + }; - const inferFd = new FormData(); - inferFd.append('session_id', sid); - inferFd.append('model_name', job.model); - inferFd.append('uploaded_filename', finalizeData.uploaded_filename || 'ct.nii.gz'); - const inferRes = await fetch(`${API_BASE}/api/run-epai-inference`, { method: 'POST', body: inferFd }); - const inferData = await parseApiResponse(inferRes); - if (!inferRes.ok) throw new Error(inferData.error || 'Failed to start inference'); - - const actualSid = inferData.session_id || sid; - if (actualSid !== sid) patchJob(job.id, { sessionId: actualSid }); - await pollUntilDone(actualSid, job.id); + const handleCheckStatus = async () => { + if (!sessionId) { setMessage("No session id yet."); return; } + try { + const res = await fetch(`${API_BASE}/api/inference-status/${sessionId}`); + const data = await parseApiResponse(res); + if (!res.ok) throw new Error(data.error || data.status || "Status check failed"); + setMessage(`Status: ${data.status}${data.error ? ` (${data.error})` : ""}`); + const status = (data.status || "").toLowerCase(); + if (status === "completed") { + setInferenceCompleted(true); + setRecentUploads(updateRecentUploadStatus(sessionId, "Completed")); + stopPolling(sessionId); + } else if (status === "cancelled") { + setRecentUploads(updateRecentUploadStatus(sessionId, "Cancelled")); + stopPolling(sessionId); + setPhase(sessionId); + } else if (status === "running" || status === "queued") { + if (!pollTimersRef.current.has(sessionId)) { + // Use the model recorded on the card - the dropdown may have changed + // since this run started, and the model decides the viewer route. + const model = loadRecentUploads().find(u => u.sessionId === sessionId)?.model || selectedModel; + startInferencePolling(sessionId, model); + } + } } catch (err) { - patchJob(job.id, { status: 'failed', error: (err as Error).message }); + console.error(err); + setMessage("Status check failed: " + (err as Error).message); } }; - /* ── Sequential queue runner ── */ - const drainQueue = async () => { - if (isProcessingRef.current) return; - isProcessingRef.current = true; - cancelledRef.current = false; - while (jobQueueRef.current.length > 0) { - if (cancelledRef.current) break; - const next = jobQueueRef.current.shift()!; - if (next.isDicom) { - await processDicomJob(next); - } else { - await processNiftiJob(next); + const handleDownloadResult = async () => { + if (!sessionId) { setMessage("No session id yet."); return; } + setMessage("Preparing download..."); + try { + const statusRes = await fetch(`${API_BASE}/api/inference-status/${sessionId}`); + const statusData = await parseApiResponse(statusRes); + if (!statusRes.ok) throw new Error(statusData.error || statusData.status || "Status check failed"); + if (statusData.status !== "completed") { + setMessage(`Status: ${statusData.status || "unknown"}. Please wait until completed.`); + return; + } + stopPolling(sessionId); + + const resultRes = await fetch(`${API_BASE}/api/get_result/${sessionId}`); + if (!resultRes.ok) { + const maybeJson = await parseApiResponse(resultRes); + throw new Error(maybeJson?.error || "Failed to download result zip"); } + const blob = await resultRes.blob(); + const objectUrl = window.URL.createObjectURL(blob); + const link = document.createElement("a"); + link.href = objectUrl; + link.download = `epai_output_${sessionId}.zip`; + document.body.appendChild(link); + link.click(); + document.body.removeChild(link); + window.URL.revokeObjectURL(objectUrl); + setMessage("Download started: zip includes combined_labels.nii.gz and output.csv"); + } catch (err) { + console.error(err); + setMessage("Download failed: " + (err as Error).message); } - isProcessingRef.current = false; }; - /* ── Run batch ── */ - const handleRun = () => { - if (!selectedModel) { alert('Please select a model first.'); return; } - if (pendingItems.length === 0) { alert('Please select at least one file.'); return; } - - const newJobs: UploadJob[] = pendingItems.map(item => ({ - id: crypto.randomUUID(), - displayName: item.displayName, - files: item.files, - isDicom: item.isDicom, - sessionId: null, - status: 'queued' as JobStatus, - uploadProgress: 0, - inferenceProgress: 0, - model: selectedModel, - timestamp: Date.now(), - })); - - setJobs(prev => [...prev, ...newJobs]); - jobQueueRef.current.push(...newJobs); - setPendingItems([]); - setPreviewItemId(null); - drainQueue(); - }; + const handleRunEpaiOnReconstruction = async () => { + if (!sessionId) { + alert("No completed reconstruction session to run ePAI on."); + return; + } + const newSessionId = crypto.randomUUID(); + setInferenceCompleted(false); + setMessage("Starting ePAI inference on reconstructed CT..."); - /* ── Cancel all inference ── */ - const handleStopAll = async () => { - cancelledRef.current = true; - jobQueueRef.current = []; - setJobs(prev => prev.map(j => - j.status === 'uploading' || j.status === 'processing' || j.status === 'queued' - ? { ...j, status: 'failed' as JobStatus, error: 'Cancelled by user' } - : j - )); - await fetch(`${API_BASE}/api/cancel-inference`, { method: 'POST' }); - }; + const formData = new FormData(); + formData.append("session_id", newSessionId); + formData.append("model_name", "ePAI"); + formData.append("source_reconstruction_session_id", sessionId); - /* ── Download ── */ - const downloadJob = async (job: UploadJob) => { - if (!job.sessionId) return; try { - const res = await fetch(`${API_BASE}/api/get_result/${job.sessionId}`); - if (!res.ok) throw new Error('Download failed'); - const blob = await res.blob(); - const url = window.URL.createObjectURL(blob); - const a = document.createElement('a'); - a.href = url; - a.download = `result_${job.sessionId}.zip`; - document.body.appendChild(a); - a.click(); - document.body.removeChild(a); - window.URL.revokeObjectURL(url); + const res = await fetch(`${API_BASE}/api/run-epai-inference`, { + method: "POST", + body: formData, + }); + const data = await parseApiResponse(res); + if (!res.ok) throw new Error(data.error || "Failed to start ePAI inference on reconstruction"); + + const sid = data.session_id || newSessionId; + setSessionId(sid); + setSelectedModel("ePAI" as const); + setMessage(`ePAI inference started on reconstructed CT. Session: ${sid}`); + if (sid) { + setRecentUploads( + addRecentUpload({ + sessionId: sid, + label: "ePAI on reconstruction", + model: "ePAI", + status: "Processing", + timestamp: Date.now(), + }) + ); + startInferencePolling(sid, "ePAI"); + } } catch (err) { - alert('Download failed: ' + (err as Error).message); + console.error(err); + setMessage("Failed to start ePAI on reconstruction: " + (err as Error).message); } }; - const handleDownloadAll = () => { - doneJobs.filter(j => j.status === 'completed').forEach(job => downloadJob(job)); - }; - - /* ── Derived state ── */ - const activeJobs = jobs.filter(j => - j.status === 'uploading' || j.status === 'processing' - ); - const queuedCount = jobs.filter(j => j.status === 'queued').length; - const activeCount = activeJobs.length + queuedCount; - const doneJobs = jobs.filter(j => j.status === 'completed' || j.status === 'failed'); - const previewItem = pendingItems.find(i => i.id === previewItemId) ?? null; - /* ── Render ── */ return (
+ {/* Ambient glow */}
@@ -478,9 +615,10 @@ const UploadPage: React.FC = () => {
Upload
- {/* Drop zone */} + {/* ── Drop zone ── */}
fileInputRef.current?.click()} onDrop={handleDrop} onDragOver={handleDragOver} onDragLeave={handleDragLeave} @@ -489,16 +627,7 @@ const UploadPage: React.FC = () => { ref={fileInputRef} type="file" multiple - accept=".nii,.gz,.dcm,.dicom" - style={{ display: 'none' }} - onChange={handleFileSelect} - /> - @@ -507,270 +636,462 @@ const UploadPage: React.FC = () => { -
Drag a file or folder here, or
-
- - -
-
.nii · .nii.gz · or a folder of .dcm slices
+
Click or drag to upload
+
.nii or .nii.gz
- {/* Pending file chips */} - {pendingItems.length > 0 && ( + {/* ── Local DICOM: view a folder of .dcm slices in-browser, nothing uploaded ── */} + )} + onChange={handleDicomFolderSelect} + /> + + + {/* ── File chips ── */} + {selectedFiles.length > 0 && (
- {pendingItems.map(item => ( -
- {item.displayName} - - + {selectedFiles.map((file, index) => ( +
+ {file.name} +
))}
)} - {/* CT Preview panel */} - {previewItem && ( -
-
Preview · {previewItem.displayName}
+ {/* ── Pre-inference preview: inspect the selected scan before running a model ── */} + {selectedFiles.length > 0 && !isUploading && ( + <> +
Preview · {selectedFiles[0].name}
Loading preview…
}> - {previewItem.isDicom - ? - : } + -
+ )} - {/* Pipeline row */} + {/* ── Pipeline row ── */}
+ {/* Step 1: Preprocessing */}
1
Preprocessing optional
- +
+ + {preDropOpen && ( +
+ {[ + { id: "", label: "None (skip)", desc: "Upload and segment as-is" }, + { id: "OpenVAE", label: "OpenVAE", desc: "Enhance the scan quality before segmenting" }, + ].map(opt => ( +
{ setPreValue(opt.id); setPreDropOpen(false); }} + > +
+ {opt.label} + {opt.desc} +
+
+ {preValue === opt.id && ( + + + + )} +
+
+ ))} +
+ )} +
→
+ {/* Step 2: Model */}
2
Model
- +
+ + {modelDropOpen && ( +
+ {MODEL_OPTIONS.map(m => ( +
{ setSelectedModel(m.id as typeof selectedModel); setModelDropOpen(false); }} + > +
+ {m.label} + {m.desc} +
+
+ {selectedModel === m.id && ( + + + + )} +
+
+ ))} +
+ )} +
→
+ {/* Step 3: Postprocessing */}
3
Postprocessing optional
- +
+ + {postDropOpen && ( +
+ {[ + { id: "", label: "None (skip)", desc: "Use results as-is" }, + { id: "ShapeKit", label: "ShapeKit", desc: "Clean up and smooth organ outlines" }, + ].map(opt => ( +
{ setPostValue(opt.id); setPostDropOpen(false); }} + > +
+ {opt.label} + {opt.desc} +
+
+ {postValue === opt.id && ( + + + + )} +
+
+ ))} +
+ )} +
-
- {/* Queue status */} - {jobs.length > 0 && ( -
- {activeCount > 0 && ( -
- + {/* ── Advanced options ── */} +
+ + {showAdvanced && ( +
+ setServerPath(e.target.value)} + /> + setBdmapId(e.target.value)} + />
)} +
- {/* Active job progress */} - {activeJobs.map(job => ( -
-
- - {job.displayName} - - {job.status === 'uploading' ? 'Uploading' : 'Running inference'} - -
+ {/* ── Action bar ── */} + {sessionId && ( +
+ + +
+ )} - {job.status === 'uploading' && ( -
- Upload -
-
-
- {job.uploadProgress}% -
- )} + {/* ── Progress (upload phase only - running inference shows in Active below) ── */} + {isUploading && ( +
+
+
+ Upload Progress + {uploadProgress}% +
+
+
+
+
+
+ )} - {job.status === 'processing' && ( + {/* ── Results ── */} + {inferenceCompleted && sessionId && ( +
+
✓ Inference Complete
+
+ {selectedModel === "OpenVAE" ? ( <> -
- Upload -
-
-
- 100% -
-
- Inference -
-
-
- {job.inferenceProgress}% -
+ + + + + ) : ( + <> + + )}
- ))} +
+ )} - {/* Queued counter */} - {activeCount > 0 && ( -
0 ? '12px' : '0' }}> - - {activeCount} scan{activeCount !== 1 ? 's' : ''} left in queue -
- )} + {/* ── Status messages ── */} + {sessionId && !inferenceCompleted && ( +
Session: {sessionId}
+ )} + {message &&
{message}
} +
- {/* Completed + failed scans */} - {doneJobs.length > 0 && ( - <> -
0 ? '20px' : '0' }}> -
- Results + {/* ── Active (in-progress) Uploads ── */} + {recentUploads.some(u => u.status === "Processing") && ( +
+
+ Active +
+
+ {recentUploads.filter(u => u.status === "Processing").map((upload) => { + const phase = sessionPhases[upload.sessionId]; + const phaseLabel = + phase === "uploading" ? "Uploading…" : + phase === "queued" ? "Queued for GPU" : + "Running…"; + return ( +
+
+
+
+
+
+
+ {upload.label} +
+
+ {upload.model ? `${upload.model} · ` : ""}{formatRelativeTime(upload.timestamp)} +
+
+
+
+ + {phaseLabel} + + +
- {doneJobs.some(j => j.status === 'completed') && ( - - )} -
+ ); + })} +
+
+ )} -
- {doneJobs.map(job => { - const isFailed = job.status === 'failed'; - return ( -
!isFailed && job.sessionId && navigate(`/session/${job.sessionId}`)} - > -
- {isFailed ? ( - - - - - - ) : ( - - - - - - - )} + {/* ── Recent Uploads ── */} +
+
+ Recent Uploads +
+
+ {recentUploads.filter(u => u.status !== "Processing").length === 0 ? ( +
+ No uploads yet - run a model above and your results will appear here. +
+ ) : ( + recentUploads.filter(u => u.status !== "Processing").map((upload) => { + const clickable = upload.status !== "Failed" && upload.status !== "Cancelled"; + const openSession = () => { + if (!clickable) return; + navigate(`/${upload.isReconstruction ? "reconstruction" : "session"}/${upload.sessionId}`); + }; + return ( +
+
+
+ + + + + + + +
+
+
+ {upload.label}
- -
-
- {job.displayName} -
-
- {isFailed - ? `Failed · ${job.error || 'Unknown error'}` - : `${job.model} · click to view`} -
+
+ {upload.model ? `${upload.model} · ` : ""}{formatRelativeTime(upload.timestamp)}
- - {!isFailed && ( - <> - - - - )}
- ); - })} -
- +
+
+ + {upload.status} + + {clickable && ( + + )} + +
+
+ ); + }) )}
- )} +
); diff --git a/PanTS-Demo/src/test/routes.smoke.test.tsx b/PanTS-Demo/src/test/routes.smoke.test.tsx index 30fdfcb..ec5933d 100644 --- a/PanTS-Demo/src/test/routes.smoke.test.tsx +++ b/PanTS-Demo/src/test/routes.smoke.test.tsx @@ -42,7 +42,7 @@ describe("route smoke tests", () => { it("UploadPage renders", async () => { renderRoute(); expect( - await screen.findByText("Drag a file or folder here, or") + await screen.findByText("Click or drag to upload") ).toBeInTheDocument(); }); }); diff --git a/flask-server/api/api_blueprint.py b/flask-server/api/api_blueprint.py index 158e5cf..42b6e1a 100644 --- a/flask-server/api/api_blueprint.py +++ b/flask-server/api/api_blueprint.py @@ -2,7 +2,7 @@ from werkzeug.utils import secure_filename from services.nifti_processor import NiftiProcessor from services.session_manager import SessionManager, generate_uuid -from services.auto_segmentor import run_auto_segmentation, cancel_all_inference +from services.auto_segmentor import run_auto_segmentation, cancel_session, cancel_all_inference from services.mesh_generation import generate_mesh_manifest, generate_organ_glb_bytes from services.inference_job_queue import InferenceJobQueue from services.intent_parser import parse_intent @@ -835,11 +835,66 @@ def download_segmentation_zip(id): _inference_jobs_lock = threading.Lock() +def _job_meta_path(session_id): + # secure_filename is the CodeQL-recognised path-injection barrier (see + # path_safety.py): session_id reaches the filesystem here, so sanitize it + # at the construction site. Real ids are UUIDs, which pass through intact. + return os.path.join(SESSIONS_DIR, secure_filename(session_id), "job.json") + + +def _persist_inference_job(session_id, snapshot): + """Mirror a job's metadata to disk so status survives a gunicorn restart. + + Best-effort and fully isolated in try/except: a disk failure here must + never break the request that triggered the status change. Atomic via + write-temp-then-rename so a crash mid-write can't leave a corrupt file. + """ + try: + path = _job_meta_path(session_id) # sanitized via secure_filename + os.makedirs(os.path.dirname(path), exist_ok=True) + tmp = path + ".tmp" + with open(tmp, "w") as f: + json.dump(snapshot, f, default=str) + os.replace(tmp, path) + except Exception as e: + print(f"[job persist] {session_id}: {e}") + + def _set_inference_job(session_id, **kwargs): with _inference_jobs_lock: current = inference_jobs.get(session_id, {}) current.update(kwargs) inference_jobs[session_id] = current + snapshot = dict(current) + _persist_inference_job(session_id, snapshot) + + +def _get_inference_job(session_id): + """Look up a job, falling back to its on-disk copy after a restart. + + Returns None when the session is genuinely unknown. A job found only on + disk was started by a *previous* process (a job from this process would + still be in memory); if that disk copy is still "running" its worker + thread died with the old process, so we surface it as failed rather than + reporting a phantom "running" that would poll forever. + """ + job = inference_jobs.get(session_id) + if job: + return job + try: + path = _job_meta_path(session_id) + if os.path.exists(path): + with open(path) as f: + disk = json.load(f) + if (disk.get("status") or "").lower() in ("running", "queued"): + disk["status"] = "failed" + disk["error"] = disk.get("error") or "Interrupted by server restart" + with _inference_jobs_lock: + inference_jobs.setdefault(session_id, disk) + return inference_jobs.get(session_id) + except Exception as e: + print(f"[job rehydrate] {session_id}: {e}") + return None def _start_auto_segmentation(session_id, model_name, ct_file=None, server_input_path=None): @@ -865,9 +920,12 @@ def _start_auto_segmentation(session_id, model_name, ct_file=None, server_input_ else: return jsonify({"error": "No CT file provided. Send MAIN_NIFTI or INPUT_SERVER_PATH."}), 400 + # "queued": the worker thread starts immediately but real GPU work waits on + # the segmentor's one-at-a-time lock; the on_start callback below flips the + # job to "running" when it actually gets the GPU. _set_inference_job( session_id, - status="running", + status="queued", model=model_name, error=None, ct_path=input_path, @@ -875,9 +933,27 @@ def _start_auto_segmentation(session_id, model_name, ct_file=None, server_input_ zip_path=os.path.join(session_path, "auto_masks.zip"), ) + def _job_status(): + job = inference_jobs.get(session_id) or {} + return (job.get("status") or "").lower() + def do_segmentation_and_zip(): + def _on_gpu_slot(): + # User may have cancelled while we sat in the queue. + if _job_status() == "cancelled": + return False + _set_inference_job(session_id, status="running") + return True + try: - output_mask_dir = run_auto_segmentation(input_path, session_dir=session_path, model=model_name) + output_mask_dir = run_auto_segmentation( + input_path, session_dir=session_path, model=model_name, + session_id=session_id, on_start=_on_gpu_slot, + ) + + if _job_status() == "cancelled": + print(f"🛑 Session {session_id} cancelled - skipping zip") + return if output_mask_dir is None or not os.path.exists(output_mask_dir): msg = f"Auto segmentation failed for session {session_id}" @@ -903,6 +979,11 @@ def do_segmentation_and_zip(): zip_path=zip_path, output_mask_dir=output_mask_dir) print(f"✅ Finished segmentation and zipping for session {session_id}") except Exception as e: + # A killed subprocess surfaces here as CalledProcessError/RuntimeError; + # if the user cancelled, keep "cancelled" rather than reporting failure. + if _job_status() == "cancelled": + print(f"🛑 Session {session_id} cancelled (worker exited: {e})") + return print(f"❌ Exception while processing session {session_id}: {e}") _set_inference_job(session_id, status="failed", error=str(e)) @@ -967,7 +1048,7 @@ def _pick_text(*keys): ) if source_reconstruction_session_id and not input_server_path and ct_file is None: - source_job = inference_jobs.get(source_reconstruction_session_id, {}) + source_job = _get_inference_job(source_reconstruction_session_id) or {} source_output_dir = source_job.get("output_mask_dir") if source_output_dir: recon_path = os.path.join(source_output_dir, "reconstructed_ct.nii.gz") @@ -1021,12 +1102,41 @@ def _pick_text(*keys): @api_blueprint.route('/inference-status/', methods=['GET']) def get_inference_status(session_id): - job = inference_jobs.get(session_id) + job = _get_inference_job(session_id) if job is None: return jsonify({"status": "not_found", "session_id": session_id}), 404 return jsonify({"session_id": session_id, **job}), 200 +@api_blueprint.route('/cancel-inference/', methods=['POST']) +def cancel_inference_session(session_id): + """Cancel one session's inference: dequeues it if still waiting for the + GPU, or SIGTERMs its subprocess group if already running. Per-session: + other users' jobs are untouched.""" + try: + if not _is_safe_id(session_id): + return jsonify({"error": "Invalid session ID"}), 400 + job = _get_inference_job(session_id) + if job is None: + return jsonify({"status": "not_found", "session_id": session_id}), 404 + status = (job.get("status") or "").lower() + if status in ("completed", "failed", "cancelled"): + return jsonify({"status": status, "message": "Job already finished"}), 200 + # Mark first so a queued job aborts at the GPU-slot check even if + # there is no subprocess to kill yet. + _set_inference_job(session_id, status="cancelled", error="Cancelled by user") + try: + killed = cancel_session(session_id) + print(f"🛑 Cancel requested for {session_id} (process killed: {killed})") + except Exception as e: + print(f"[cancel] kill failed for {session_id}: {e}") + return jsonify({"status": "cancelled", "session_id": session_id}), 200 + except Exception as e: + # Log the detail server-side; don't leak internals to the client. + print(f"❌ [Cancel Error] {e}") + return jsonify({"error": "Failed to cancel inference"}), 500 + + @api_blueprint.route('/check-inference-status', methods=['GET']) def check_inference_status_legacy(): session_id = request.args.get("session_id") or request.args.get("sessionId") @@ -1266,7 +1376,7 @@ def get_result(session_id): @api_blueprint.route('/session-ct/', methods=['GET']) def get_session_ct(session_id): - job = inference_jobs.get(session_id, {}) + job = _get_inference_job(session_id) or {} ct_path = job.get("ct_path") if not ct_path or not os.path.exists(ct_path): return jsonify({"error": "CT file not found for session"}), 404 @@ -1278,7 +1388,7 @@ def get_session_ct(session_id): @api_blueprint.route('/session-segmentation/', methods=['GET']) def get_session_segmentation(session_id): - job = inference_jobs.get(session_id, {}) + job = _get_inference_job(session_id) or {} output_mask_dir = job.get("output_mask_dir") if not output_mask_dir: return jsonify({"error": "Segmentation not ready for session"}), 404 @@ -1294,7 +1404,7 @@ def get_session_segmentation(session_id): @api_blueprint.route('/session-reconstruction/', methods=['GET']) def get_session_reconstruction(session_id): """Serves the OpenVAE reconstructed CT for a session.""" - job = inference_jobs.get(session_id, {}) + job = _get_inference_job(session_id) or {} output_mask_dir = job.get("output_mask_dir") if not output_mask_dir: return jsonify({"error": "Reconstruction not ready for session"}), 404 diff --git a/flask-server/services/auto_segmentor.py b/flask-server/services/auto_segmentor.py index 833188d..6a71ce0 100644 --- a/flask-server/services/auto_segmentor.py +++ b/flask-server/services/auto_segmentor.py @@ -1,7 +1,7 @@ import os import uuid -import subprocess import signal +import subprocess import re import csv import shlex @@ -15,44 +15,85 @@ # Only one model inference runs at a time to avoid GPU OOM _gpu_lock = threading.Lock() -# Track the active subprocess so it can be killed -_active_proc: 'subprocess.Popen | None' = None -_active_proc_lock = threading.Lock() +# ── Per-session subprocess tracking (for user-initiated cancel) ── +# Each session's worker thread binds its session id thread-locally; _tracked_run +# then registers the live Popen under that id so /api/cancel-inference/ +# can kill exactly that session's process group and nobody else's. +_session_procs = {} +_session_procs_lock = threading.Lock() +_thread_session = threading.local() -def cancel_all_inference(): - """Kill the currently running inference subprocess (if any).""" - global _active_proc - with _active_proc_lock: - proc = _active_proc - if proc and proc.poll() is None: - try: - os.killpg(os.getpgid(proc.pid), signal.SIGTERM) - except Exception: +def bind_session(session_id): + """Associate the calling worker thread with a session id.""" + _thread_session.sid = session_id + + +def _tracked_run(cmd, check=False, capture_output=False, **kwargs): + """Drop-in for subprocess.run that makes the child killable per-session. + + Starts the child in its own process group (start_new_session) and registers + it under the thread's bound session id for cancel_session(). Mirrors the + subprocess.run semantics used in this module (check / capture_output / + shell / executable / cwd / stdout / stderr / text). + """ + if capture_output: + kwargs.setdefault("stdout", subprocess.PIPE) + kwargs.setdefault("stderr", subprocess.PIPE) + kwargs.setdefault("start_new_session", True) + sid = getattr(_thread_session, "sid", None) + proc = subprocess.Popen(cmd, **kwargs) + if sid: + with _session_procs_lock: + _session_procs[sid] = proc + try: + stdout, stderr = proc.communicate() + finally: + if sid: + with _session_procs_lock: + if _session_procs.get(sid) is proc: + _session_procs.pop(sid, None) + if check and proc.returncode != 0: + raise subprocess.CalledProcessError(proc.returncode, cmd, output=stdout, stderr=stderr) + return subprocess.CompletedProcess(cmd, proc.returncode, stdout, stderr) + + +def cancel_session(session_id): + """Best-effort kill of one session's inference subprocess (SIGTERM the + process group, SIGKILL 5s later if it ignores that). Returns True if a + live process was signalled. Safe to call for queued/unknown sessions.""" + with _session_procs_lock: + proc = _session_procs.get(session_id) + if not proc or proc.poll() is not None: + return False + try: + pgid = os.getpgid(proc.pid) + os.killpg(pgid, signal.SIGTERM) + + def _force_kill(): try: - proc.kill() + proc.wait(timeout=5) except Exception: - pass - - -def _run_subprocess(cmd: str, **kwargs): - """Run a shell command, tracking it so cancel_all_inference() can kill it.""" - global _active_proc - proc = subprocess.Popen( - cmd, - shell=True, - executable='/bin/bash', - start_new_session=True, - **kwargs, - ) - with _active_proc_lock: - _active_proc = proc - ret = proc.wait() - with _active_proc_lock: - if _active_proc is proc: - _active_proc = None - if ret != 0: - raise subprocess.CalledProcessError(ret, cmd) + try: + os.killpg(pgid, signal.SIGKILL) + except Exception: + pass + + threading.Thread(target=_force_kill, daemon=True).start() + return True + except Exception as e: + print(f"[cancel] failed to kill process for session {session_id}: {e}") + return False + + +def cancel_all_inference(): + """Kill every tracked inference subprocess (admin 'stop everything'). + Prefer cancel_session() for a single user's job; this is the blunt + kill-all kept for the global /cancel-inference endpoint.""" + with _session_procs_lock: + sids = list(_session_procs.keys()) + for sid in sids: + cancel_session(sid) def get_least_used_gpu(default_gpu=None): if default_gpu is None: @@ -92,13 +133,27 @@ def _resolve_conda_activate_path(): return "" -def run_auto_segmentation(input_path, session_dir, model): +def run_auto_segmentation(input_path, session_dir, model, session_id=None, on_start=None): """ Dispatch to the appropriate model inference function. Serialized via _gpu_lock so concurrent requests queue instead of OOM-ing. Returns the output directory path on success, raises on failure. + + session_id: binds this worker thread so _tracked_run/cancel_session can + target its subprocesses. on_start: called once the GPU slot is + acquired (i.e. the job leaves the queue); returning False aborts the + run (used when the user cancelled while the job was still queued) and + makes this function return None. """ with _gpu_lock: + if on_start is not None: + try: + if on_start() is False: + return None + except Exception as e: + print(f"[on_start] callback error for {session_id}: {e}") + if session_id: + bind_session(session_id) if model == 'ePAI': conda_path = _resolve_conda_activate_path() return _run_epai_inference( @@ -403,7 +458,12 @@ def _run_epai_inference(input_path: str, session_dir: str, conda_path: str, epai print(f"[INFO] Running ePAI command for case {case_id}") print(full_cmd) try: - _run_subprocess(full_cmd) + _tracked_run( + full_cmd, + shell=True, + executable="/bin/bash", + check=True, + ) except subprocess.CalledProcessError as e: raise RuntimeError( "ePAI inference command failed" @@ -473,7 +533,8 @@ def _run_suprem_inference(input_path: str, session_dir: str) -> str: print(f"[INFO] Running SuPreM native inference") print(full_cmd) try: - _run_subprocess(full_cmd, cwd=suprem_src) + _tracked_run(full_cmd, shell=True, executable="/bin/bash", check=True, + cwd=suprem_src) except subprocess.CalledProcessError as e: raise RuntimeError( f"SuPreM inference failed\nCommand: {full_cmd}\nExit code: {e.returncode}" @@ -532,7 +593,7 @@ def _run_openvae_inference(input_path: str, session_dir: str) -> str: ) print(f"[INFO] Running OpenVAE inference\n{full_cmd}") try: - _run_subprocess(full_cmd, cwd=openvae_src) + _tracked_run(full_cmd, shell=True, executable="/bin/bash", check=True, cwd=openvae_src) except subprocess.CalledProcessError as e: raise RuntimeError( f"OpenVAE inference failed\nCommand: {full_cmd}\nExit code: {e.returncode}" @@ -624,7 +685,10 @@ def _run_medformer_inference(input_path: str, session_dir: str) -> str: ) print(f"[INFO] Running MedFormer inference\n{full_cmd}") try: - _run_subprocess(full_cmd, cwd=rsuper_src) + _tracked_run( + full_cmd, shell=True, executable="/bin/bash", check=True, cwd=rsuper_src, + stdout=subprocess.PIPE, stderr=subprocess.STDOUT + ) except subprocess.CalledProcessError as e: raise RuntimeError( f"MedFormer inference failed\nCommand: {full_cmd}\nExit code: {e.returncode}" @@ -684,7 +748,10 @@ def _run_rsuper_inference(input_path: str, session_dir: str) -> str: ) print(f"[INFO] Running R-Super inference\n{full_cmd}") try: - _run_subprocess(full_cmd, cwd=rsuper_src) + _tracked_run( + full_cmd, shell=True, executable="/bin/bash", check=True, cwd=rsuper_src, + stdout=subprocess.PIPE, stderr=subprocess.STDOUT + ) except subprocess.CalledProcessError as e: raise RuntimeError( f"R-Super inference failed\nCommand: {full_cmd}\nExit code: {e.returncode}" @@ -744,7 +811,7 @@ def _run_atlasnet_inference(input_path: str, session_dir: str, conda_path: str, print(f"[INFO] Running Atlas-Net command for case {case_id}") print(full_cmd) try: - _run_subprocess(full_cmd) + _tracked_run(full_cmd, shell=True, executable="/bin/bash", check=True) except subprocess.CalledProcessError as e: raise RuntimeError( f"Atlas-Net inference command failed\nCommand: {full_cmd}\nExit code: {e.returncode}" @@ -768,7 +835,7 @@ def _is_truthy(value: str) -> bool: def _run_checked_process(cmd: list[str], error_prefix: str): - process = subprocess.run(cmd, text=True, capture_output=True) + process = _tracked_run(cmd, text=True, capture_output=True) if process.returncode != 0: raise RuntimeError( f"{error_prefix}" @@ -983,7 +1050,7 @@ def _run_shapekit_inference(input_dir: str, session_dir: str) -> str: ) print(f"[INFO] Running ShapeKit post-processing\n{full_cmd}") try: - _run_subprocess(full_cmd, cwd=shapekit_src) + _tracked_run(full_cmd, shell=True, executable="/bin/bash", check=True, cwd=shapekit_src) except subprocess.CalledProcessError as e: raise RuntimeError( f"ShapeKit post-processing failed\nCommand: {full_cmd}\nExit code: {e.returncode}"