From eac39cd196da3c345abb8e60ab3419417d8973fe Mon Sep 17 00:00:00 2001 From: Aditya Sanjeev Date: Wed, 15 Jul 2026 00:48:01 -0700 Subject: [PATCH 1/4] Upload page: resumable/parallel uploads, per-session cancel, disk-backed jobs Frontend (UploadPage): - Resumable uploads: the picked file is stashed in IndexedDB with a chunk cursor, so a tab close mid-upload resumes on reopen instead of dying. - Parallel runs with independent per-session polling; an "Active" section shows each in-flight run with a phase label (Uploading / Queued for GPU / Running) and a per-card Cancel button. - Testing model: a client-side 20s dry run for local dev without a GPU. - Delete button and a terminal "Cancelled" status on Recent Uploads. Backend: - Per-session subprocess tracking (_tracked_run / bind_session / cancel_session) replaces the global single-process tracking, so /api/cancel-inference/ kills exactly one session's process group and leaves other users' jobs running. The global /cancel-inference kill-all is kept and re-implemented on the same registry. - inference_jobs is mirrored to tmp//job.json (atomic write) and read back through _get_inference_job, so status/viewer endpoints survive a gunicorn restart; a job left "running"/"queued" on disk by a dead process is coerced to "failed" instead of polling forever. - Job lifecycle is queued -> running -> completed/failed/cancelled; an on_start callback flips queued->running at GPU acquisition and aborts a run cancelled while still queued. Renamed the per-session cancel view to cancel_inference_session to avoid a duplicate Flask endpoint name with the existing global cancel view. --- PanTS-Demo/src/helpers/pendingUploads.ts | 106 ++ PanTS-Demo/src/helpers/recentUploads.ts | 16 +- PanTS-Demo/src/routes/UploadPage.css | 129 ++ PanTS-Demo/src/routes/UploadPage.tsx | 1575 +++++++++++++-------- PanTS-Demo/src/test/routes.smoke.test.tsx | 2 +- flask-server/api/api_blueprint.py | 123 +- flask-server/services/auto_segmentor.py | 155 +- 7 files changed, 1458 insertions(+), 648 deletions(-) create mode 100644 PanTS-Demo/src/helpers/pendingUploads.ts diff --git a/PanTS-Demo/src/helpers/pendingUploads.ts b/PanTS-Demo/src/helpers/pendingUploads.ts new file mode 100644 index 0000000..70a4d42 --- /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..e9d023c 100644 --- a/PanTS-Demo/src/routes/UploadPage.tsx +++ b/PanTS-Demo/src/routes/UploadPage.tsx @@ -1,131 +1,66 @@ 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" }, + { id: "Testing", label: "Testing", desc: "Local 20-second dry run" }, +]; 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.', +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 TESTING_LS_KEY = "uploadTestingSessions"; +const TESTING_DURATION_MS = 20_000; + +// sessionId → startedAt (ms). A map so several Testing runs can be in flight. +const loadTestingSessions = (): Record => { + try { + return JSON.parse(localStorage.getItem(TESTING_LS_KEY) || "{}") || {}; + } catch { return {}; } }; -type JobStatus = 'queued' | 'uploading' | 'processing' | 'completed' | 'failed'; - -type PendingItem = { - id: string; - displayName: string; - files: File[]; - isDicom: boolean; -}; +const saveTestingSession = (sessionId: string, startedAt: number) => + localStorage.setItem(TESTING_LS_KEY, JSON.stringify({ ...loadTestingSessions(), [sessionId]: startedAt })); -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; +const clearTestingSession = (sessionId: string) => { + const sessions = loadTestingSessions(); + delete sessions[sessionId]; + localStorage.setItem(TESTING_LS_KEY, JSON.stringify(sessions)); }; 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 +68,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" | "Testing" | "">(""); + 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 +159,519 @@ const UploadPage: React.FC = () => { setIsDragOver(true); }, []); - const handleDragLeave = useCallback(() => setIsDragOver(false), []); + const handleDragLeave = useCallback(() => { + setIsDragOver(false); + }, []); + + 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 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."); + } + 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); + }; + + const startTestingTimer = (sid: string, startedAt: number) => { + stopPolling(sid); + const timer = setInterval(() => { + if (Date.now() - startedAt >= TESTING_DURATION_MS) { + clearTestingSession(sid); + finishSession(sid, "Testing"); + } + }, 200); + pollTimersRef.current.set(sid, timer); + }; + + // Cancel one run, whatever phase it's in: aborts an in-flight upload, kills + // a queued/running server job, or just stops a Testing dry run. + const cancelRun = (upload: RecentUpload) => { + const sid = upload.sessionId; + stopPolling(sid); + setPhase(sid); + + const controller = uploadAbortRef.current.get(sid); + if (controller) controller.abort(); + deletePendingUpload(sid); + + if (upload.model === "Testing") { + clearTestingSession(sid); + } else { + // 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(() => {}); + } - const removePendingItem = (id: string) => { - setPendingItems(prev => prev.filter(item => item.id !== id)); - if (previewItemId === id) setPreviewItemId(null); + if (foregroundUploadSidRef.current === sid) { + foregroundUploadSidRef.current = null; + setIsUploading(false); + } + setRecentUploads(updateRecentUploadStatus(sid, "Cancelled")); + setMessage(`Cancelled ${upload.label}`); }; - 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(); + 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 testingSessions = loadTestingSessions(); + const pending = await loadPendingUploads(); + if (cancelled) return; + const pendingById = new Map(pending.map(p => [p.sessionId, p])); + + for (const u of processing) { + if (u.model === "Testing") { + const startedAt = testingSessions[u.sessionId]; + if (startedAt == null) { + setRecentUploads(updateRecentUploadStatus(u.sessionId, "Failed")); + } else if (Date.now() - startedAt >= TESTING_DURATION_MS) { + clearTestingSession(u.sessionId); + finishSession(u.sessionId, "Testing"); } else { - setJobs(prev => prev.map(j => - j.id === jobId - ? { ...j, inferenceProgress: Math.min(95, j.inferenceProgress + 7) } - : j - )); + startTestingTimer(u.sessionId, startedAt); } - } catch { /* keep polling */ } - }, 2500); - }); + } else if (pendingById.has(u.sessionId)) { + runUpload(pendingById.get(u.sessionId)!, false); // resume the upload + } else { + startInferencePolling(u.sessionId, u.model); // resume polling + } + } + + // Clean up storage entries whose card no longer exists (deleted or + // trimmed off the 8-entry list) so neither store can leak. + const known = new Set(loadRecentUploads().map(u => u.sessionId)); + Object.keys(testingSessions).filter(sid => !known.has(sid)).forEach(clearTestingSession); + pending.filter(p => !known.has(p.sessionId)).forEach(p => deletePendingUpload(p.sessionId)); - /* ── Persistence: mirror the queue to localStorage on every change ── */ + 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 () => { + if (selectedModel === "Testing") { + const sid = crypto.randomUUID(); + const startedAt = Date.now(); + setSessionId(sid); + setMessage("Testing mode — 20-second dry run"); + setInferenceCompleted(false); + setRecentUploads( + addRecentUpload({ + sessionId: sid, + label: selectedFiles[0]?.name || "Testing run", + model: "Testing", + status: "Processing", + timestamp: startedAt, + }) + ); + if (selectedFiles.length > 0) setSelectedFiles(prev => prev.slice(1)); + saveTestingSession(sid, startedAt); + startTestingTimer(sid, startedAt); + return; + } + + 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 +683,10 @@ const UploadPage: React.FC = () => {
Upload
- {/* Drop zone */} + {/* ── Drop zone ── */}
fileInputRef.current?.click()} onDrop={handleDrop} onDragOver={handleDragOver} onDragLeave={handleDragLeave} @@ -489,16 +695,7 @@ const UploadPage: React.FC = () => { ref={fileInputRef} type="file" multiple - accept=".nii,.gz,.dcm,.dicom" - style={{ display: 'none' }} - onChange={handleFileSelect} - /> - @@ -507,270 +704,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..c1897af 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,64 @@ def download_segmentation_zip(id): _inference_jobs_lock = threading.Lock() +def _job_meta_path(session_id): + return os.path.join(SESSIONS_DIR, 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: + session_dir = os.path.join(SESSIONS_DIR, session_id) + os.makedirs(session_dir, exist_ok=True) + path = _job_meta_path(session_id) + 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 +918,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 +931,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 +977,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 +1046,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 +1100,40 @@ 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: + print(f"❌ [Cancel Error] {e}") + return jsonify({"error": str(e)}), 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 +1373,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 +1385,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 +1401,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}" From 466e88ddef64fcf4370230abc90303493d585b73 Mon Sep 17 00:00:00 2001 From: Aditya Sanjeev Date: Wed, 15 Jul 2026 00:52:32 -0700 Subject: [PATCH 2/4] Drop the local-only Testing model from the upload page It was dev scaffolding for exercising the flow without a GPU; not meant to ship. Removes the Testing option, its localStorage timer state, and the resume/cancel branches that special-cased it. --- PanTS-Demo/src/routes/UploadPage.tsx | 86 +++------------------------- 1 file changed, 9 insertions(+), 77 deletions(-) diff --git a/PanTS-Demo/src/routes/UploadPage.tsx b/PanTS-Demo/src/routes/UploadPage.tsx index e9d023c..32e552a 100644 --- a/PanTS-Demo/src/routes/UploadPage.tsx +++ b/PanTS-Demo/src/routes/UploadPage.tsx @@ -6,7 +6,6 @@ const MODEL_OPTIONS: { id: string; label: string; desc: string }[] = [ { 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" }, - { id: "Testing", label: "Testing", desc: "Local 20-second dry run" }, ]; import { useNavigate } from 'react-router-dom'; import './UploadPage.css'; @@ -33,25 +32,6 @@ import { type PendingUpload, } from '../helpers/pendingUploads'; -const TESTING_LS_KEY = "uploadTestingSessions"; -const TESTING_DURATION_MS = 20_000; - -// sessionId → startedAt (ms). A map so several Testing runs can be in flight. -const loadTestingSessions = (): Record => { - try { - return JSON.parse(localStorage.getItem(TESTING_LS_KEY) || "{}") || {}; - } catch { return {}; } -}; - -const saveTestingSession = (sessionId: string, startedAt: number) => - localStorage.setItem(TESTING_LS_KEY, JSON.stringify({ ...loadTestingSessions(), [sessionId]: startedAt })); - -const clearTestingSession = (sessionId: string) => { - const sessions = loadTestingSessions(); - delete sessions[sessionId]; - localStorage.setItem(TESTING_LS_KEY, JSON.stringify(sessions)); -}; - const parseApiResponse = async (res: Response): Promise => { const contentType = res.headers.get("content-type") || ""; if (contentType.includes("application/json")) { @@ -100,7 +80,7 @@ const UploadPage: React.FC = () => { 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" | "Testing" | "">(""); + 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); @@ -246,19 +226,8 @@ const UploadPage: React.FC = () => { pollTimersRef.current.set(sid, timer); }; - const startTestingTimer = (sid: string, startedAt: number) => { - stopPolling(sid); - const timer = setInterval(() => { - if (Date.now() - startedAt >= TESTING_DURATION_MS) { - clearTestingSession(sid); - finishSession(sid, "Testing"); - } - }, 200); - pollTimersRef.current.set(sid, timer); - }; - - // Cancel one run, whatever phase it's in: aborts an in-flight upload, kills - // a queued/running server job, or just stops a Testing dry run. + // 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); @@ -268,13 +237,9 @@ const UploadPage: React.FC = () => { if (controller) controller.abort(); deletePendingUpload(sid); - if (upload.model === "Testing") { - clearTestingSession(sid); - } else { - // 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(() => {}); - } + // 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; @@ -292,33 +257,21 @@ const UploadPage: React.FC = () => { let cancelled = false; (async () => { const processing = loadRecentUploads().filter(u => u.status === "Processing"); - const testingSessions = loadTestingSessions(); const pending = await loadPendingUploads(); if (cancelled) return; const pendingById = new Map(pending.map(p => [p.sessionId, p])); for (const u of processing) { - if (u.model === "Testing") { - const startedAt = testingSessions[u.sessionId]; - if (startedAt == null) { - setRecentUploads(updateRecentUploadStatus(u.sessionId, "Failed")); - } else if (Date.now() - startedAt >= TESTING_DURATION_MS) { - clearTestingSession(u.sessionId); - finishSession(u.sessionId, "Testing"); - } else { - startTestingTimer(u.sessionId, startedAt); - } - } else if (pendingById.has(u.sessionId)) { + if (pendingById.has(u.sessionId)) { runUpload(pendingById.get(u.sessionId)!, false); // resume the upload } else { startInferencePolling(u.sessionId, u.model); // resume polling } } - // Clean up storage entries whose card no longer exists (deleted or - // trimmed off the 8-entry list) so neither store can leak. + // 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)); - Object.keys(testingSessions).filter(sid => !known.has(sid)).forEach(clearTestingSession); pending.filter(p => !known.has(p.sessionId)).forEach(p => deletePendingUpload(p.sessionId)); if (processing.length > 0) { @@ -476,27 +429,6 @@ const UploadPage: React.FC = () => { /* ── Run inference ── */ const handleRunEpaiInference = async () => { - if (selectedModel === "Testing") { - const sid = crypto.randomUUID(); - const startedAt = Date.now(); - setSessionId(sid); - setMessage("Testing mode — 20-second dry run"); - setInferenceCompleted(false); - setRecentUploads( - addRecentUpload({ - sessionId: sid, - label: selectedFiles[0]?.name || "Testing run", - model: "Testing", - status: "Processing", - timestamp: startedAt, - }) - ); - if (selectedFiles.length > 0) setSelectedFiles(prev => prev.slice(1)); - saveTestingSession(sid, startedAt); - startTestingTimer(sid, startedAt); - return; - } - const file = selectedFiles[0] ?? null; const path = serverPath.trim(); if (!file && !path) { From c1dd119e9e66c6bd4308e825de6a9cdc95826af3 Mon Sep 17 00:00:00 2001 From: Aditya Sanjeev Date: Wed, 15 Jul 2026 01:01:04 -0700 Subject: [PATCH 3/4] Harden disk-backed job I/O against CodeQL path-injection + stack-trace exposure - Route job.json paths through secure_filename (the CodeQL-recognised path-injection barrier per path_safety.py); session_id reaches the filesystem in _job_meta_path/_persist_inference_job. Real ids are UUIDs and pass through unchanged. - Per-session cancel endpoint no longer returns str(e) to the client; logs the detail server-side and returns a generic message. --- flask-server/api/api_blueprint.py | 13 ++++++++----- 1 file changed, 8 insertions(+), 5 deletions(-) diff --git a/flask-server/api/api_blueprint.py b/flask-server/api/api_blueprint.py index c1897af..16be036 100644 --- a/flask-server/api/api_blueprint.py +++ b/flask-server/api/api_blueprint.py @@ -836,7 +836,10 @@ def download_segmentation_zip(id): def _job_meta_path(session_id): - return os.path.join(SESSIONS_DIR, session_id, "job.json") + # 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): @@ -847,9 +850,8 @@ def _persist_inference_job(session_id, snapshot): write-temp-then-rename so a crash mid-write can't leave a corrupt file. """ try: - session_dir = os.path.join(SESSIONS_DIR, session_id) - os.makedirs(session_dir, exist_ok=True) - path = _job_meta_path(session_id) + 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) @@ -1130,8 +1132,9 @@ def cancel_inference_session(session_id): 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": str(e)}), 500 + return jsonify({"error": "Failed to cancel inference"}), 500 @api_blueprint.route('/check-inference-status', methods=['GET']) From e4a4a04f5ff6e4118baf786fad758193649285de Mon Sep 17 00:00:00 2001 From: Aditya Sanjeev Date: Wed, 15 Jul 2026 01:11:21 -0700 Subject: [PATCH 4/4] Replace em dashes with hyphens in upload flow strings and comments --- PanTS-Demo/src/helpers/pendingUploads.ts | 6 ++--- PanTS-Demo/src/routes/UploadPage.tsx | 30 ++++++++++++------------ flask-server/api/api_blueprint.py | 4 ++-- 3 files changed, 20 insertions(+), 20 deletions(-) diff --git a/PanTS-Demo/src/helpers/pendingUploads.ts b/PanTS-Demo/src/helpers/pendingUploads.ts index 70a4d42..1b61600 100644 --- a/PanTS-Demo/src/helpers/pendingUploads.ts +++ b/PanTS-Demo/src/helpers/pendingUploads.ts @@ -2,7 +2,7 @@ // // 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 — +// 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. @@ -12,7 +12,7 @@ export type PendingUpload = { sessionId: string; - file: Blob; // the picked File — File extends Blob, survives in IDB + file: Blob; // the picked File - File extends Blob, survives in IDB filename: string; model: string; bdmapId: string; @@ -59,7 +59,7 @@ export async function savePendingUpload(p: PendingUpload): Promise { await withStore("readwrite", store => store.put(p)); return true; } catch (e) { - console.warn("savePendingUpload failed — upload won't be resumable", e); + console.warn("savePendingUpload failed - upload won't be resumable", e); return false; } } diff --git a/PanTS-Demo/src/routes/UploadPage.tsx b/PanTS-Demo/src/routes/UploadPage.tsx index 32e552a..1f1b528 100644 --- a/PanTS-Demo/src/routes/UploadPage.tsx +++ b/PanTS-Demo/src/routes/UploadPage.tsx @@ -59,7 +59,7 @@ const UploadPage: React.FC = () => { 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. + // 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 @@ -194,7 +194,7 @@ const UploadPage: React.FC = () => { stopPolling(sid); setPhase(sid); setRecentUploads(updateRecentUploadStatus(sid, "Failed")); - setMessage("Session no longer exists on the server — marked as Failed."); + setMessage("Session no longer exists on the server - marked as Failed."); } return; } @@ -210,7 +210,7 @@ const UploadPage: React.FC = () => { setRecentUploads(updateRecentUploadStatus(sid, "Failed")); setMessage(`Inference failed${data.error ? `: ${data.error}` : ""}`); } else if (status === "cancelled") { - // Cancelled elsewhere (another tab, or the backend) — reflect it. + // Cancelled elsewhere (another tab, or the backend) - reflect it. stopPolling(sid); setPhase(sid); setRecentUploads(updateRecentUploadStatus(sid, "Cancelled")); @@ -218,7 +218,7 @@ const UploadPage: React.FC = () => { setPhase(sid, status); } } catch (err) { - // Network blip or proxy error while the backend restarts — the job + // Network blip or proxy error while the backend restarts - the job // may still be alive server-side, so keep polling. console.error(err); } @@ -238,7 +238,7 @@ const UploadPage: React.FC = () => { 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. + // 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) { @@ -250,9 +250,9 @@ const UploadPage: React.FC = () => { }; useEffect(() => { - // Resume every in-flight run — there can be several in parallel. Uploads + // 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 + // (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 () => { @@ -283,7 +283,7 @@ const UploadPage: React.FC = () => { }, []); // Only warn before an unload if the current upload could NOT be stored in - // IndexedDB (quota/private-mode) — otherwise an interrupted upload resumes + // IndexedDB (quota/private-mode) - otherwise an interrupted upload resumes // automatically on reopen, so no scary dialog is needed. useEffect(() => { if (!isUploading || uploadResumableRef.current) return; @@ -361,7 +361,7 @@ const UploadPage: React.FC = () => { const data = await parseApiResponse(res); if (!res.ok) throw new Error(data.error || "Chunk upload failed"); - // Persist the cursor every so often (not every chunk — that would be a + // 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); @@ -383,7 +383,7 @@ const UploadPage: React.FC = () => { 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 + // 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) { @@ -408,7 +408,7 @@ const UploadPage: React.FC = () => { if (foreground) setMessage(`${model} inference started. Session: ${sid}`); startInferencePolling(sid, model); } catch (err) { - // A user cancel aborts our fetches — cancelRun already did the cleanup + // 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); @@ -510,7 +510,7 @@ const UploadPage: React.FC = () => { 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 + // 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); @@ -652,7 +652,7 @@ const UploadPage: React.FC = () => { /> {/* ── File chips ── */} @@ -868,7 +868,7 @@ const UploadPage: React.FC = () => {
)} - {/* ── Progress (upload phase only — running inference shows in Active below) ── */} + {/* ── Progress (upload phase only - running inference shows in Active below) ── */} {isUploading && (
@@ -1011,7 +1011,7 @@ const UploadPage: React.FC = () => { fontSize: "12px", color: "#8f8f8f" }}> - No uploads yet — run a model above and your results will appear here. + No uploads yet - run a model above and your results will appear here.
) : ( recentUploads.filter(u => u.status !== "Processing").map((upload) => { diff --git a/flask-server/api/api_blueprint.py b/flask-server/api/api_blueprint.py index 16be036..42b6e1a 100644 --- a/flask-server/api/api_blueprint.py +++ b/flask-server/api/api_blueprint.py @@ -952,7 +952,7 @@ def _on_gpu_slot(): ) if _job_status() == "cancelled": - print(f"🛑 Session {session_id} cancelled — skipping zip") + print(f"🛑 Session {session_id} cancelled - skipping zip") return if output_mask_dir is None or not os.path.exists(output_mask_dir): @@ -1111,7 +1111,7 @@ def get_inference_status(session_id): @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 — + GPU, or SIGTERMs its subprocess group if already running. Per-session: other users' jobs are untouched.""" try: if not _is_safe_id(session_id):