diff --git a/PanTS-Demo/src/routes/UploadPage.css b/PanTS-Demo/src/routes/UploadPage.css index 2841af3..41e5888 100644 --- a/PanTS-Demo/src/routes/UploadPage.css +++ b/PanTS-Demo/src/routes/UploadPage.css @@ -115,6 +115,52 @@ letter-spacing: 0.04em; } +/* ── Dropzone buttons ── */ +.dropzone-btn { + padding: 7px 16px; + background: rgba(0,0,0,.06); + border: 1px solid rgba(0,0,0,.12); + border-radius: 8px; + font-family: 'Space Grotesk', sans-serif; + font-size: 12px; + font-weight: 500; + color: #111111; + cursor: pointer; + transition: all 0.15s; +} +.dropzone-btn:hover { + background: rgba(0,0,0,.10); + border-color: rgba(0,0,0,.2); +} + +/* ── Local DICOM opener under the dropzone ── */ +.dicom-open-link { + display: flex; + flex-direction: column; + align-items: center; + gap: 2px; + width: 100%; + margin-top: 10px; + padding: 10px 12px; + background: transparent; + border: 1px dashed rgba(0, 0, 0, 0.18); + border-radius: 10px; + color: rgba(0, 0, 0, 0.55); + font-size: 13px; + cursor: pointer; + transition: border-color 0.15s, color 0.15s; +} +.dicom-open-link:hover { + border-color: rgba(0, 0, 0, 0.35); + color: #111111; +} +.dicom-open-link span { + font-family: 'JetBrains Mono', monospace; + font-size: 10.5px; + color: #8f8f8f; + letter-spacing: 0.04em; +} + /* ── File chips ── */ .file-chips { display: flex; @@ -205,6 +251,170 @@ letter-spacing: 0.02em; text-transform: none; } +/* ── Custom pipeline select ── */ +.ps-wrap { + position: relative; + width: 100%; +} + +.ps-trigger { + width: 100%; + display: flex; + align-items: center; + justify-content: space-between; + padding: 10px 12px 10px 14px; + background: rgba(0,0,0,.04); + border: 1px solid rgba(0,0,0,.08); + border-radius: 8px; + font-family: 'Space Grotesk', sans-serif; + font-size: 13px; + color: #8f8f8f; + cursor: pointer; + text-align: left; + transition: all 0.2s; + outline: none; +} + +.ps-trigger.has-value { + color: #111; +} + +.ps-trigger:hover { + border-color: rgba(0,0,0,.15); + background-color: rgba(0,0,0,.06); +} + +.ps-arrow { + flex-shrink: 0; + margin-left: 8px; + opacity: 0; + width: 0; +} + +.ps-trigger::after { + content: ''; + display: block; + width: 10px; + height: 6px; + flex-shrink: 0; + background-image: url("data:image/svg+xml,%3Csvg width='10' height='6' viewBox='0 0 10 6' fill='none' xmlns='http://www.w3.org/2000/svg'%3E%3Cpath d='M1 1l4 4 4-4' stroke='%236a6a6a' stroke-width='1.5' stroke-linecap='round' stroke-linejoin='round'/%3E%3C/svg%3E"); + background-repeat: no-repeat; + background-position: center; +} + +.ps-dropdown { + position: absolute; + top: calc(100% + 4px); + left: 0; + right: 0; + background: #fff; + border: 1px solid rgba(0,0,0,.12); + border-radius: 8px; + box-shadow: 0 4px 16px rgba(0,0,0,0.10); + z-index: 200; + overflow: visible; +} + +.ps-option { + display: flex; + align-items: center; + justify-content: space-between; + padding: 9px 12px; + font-family: 'Space Grotesk', sans-serif; + font-size: 13px; + color: #111; + cursor: pointer; + transition: background 0.12s; +} + +.ps-option:first-child { border-radius: 8px 8px 0 0; } +.ps-option:last-child { border-radius: 0 0 8px 8px; } +.ps-option:only-child { border-radius: 8px; } + +.ps-option:hover { + background: rgba(0,0,0,.05); +} + +.ps-option--selected { + background: rgba(0,0,0,.07); + color: #111; + font-weight: 500; +} + +.ps-option--selected:hover { + background: rgba(0,0,0,.10); +} + +.ps-option-label { + flex: 1; +} + +/* ── Info icon & tooltip ── */ +.ps-info-wrap { + position: relative; + display: inline-flex; + align-items: center; + margin-left: 8px; + flex-shrink: 0; +} + +.ps-info-icon { + display: inline-flex; + align-items: center; + justify-content: center; + width: 15px; + height: 15px; + border-radius: 50%; + border: 1px solid currentColor; + font-size: 9px; + font-style: italic; + font-family: Georgia, serif; + font-weight: 700; + cursor: default; + opacity: 0.5; + transition: opacity 0.15s; +} + +.ps-info-wrap:hover .ps-info-icon { + opacity: 1; +} + + +.ps-info-tooltip { + display: none; + position: absolute; + right: calc(100% + 8px); + top: 50%; + transform: translateY(-50%); + width: 210px; + background: #111; + color: #f5f5f5; + font-family: 'Space Grotesk', sans-serif; + font-size: 11px; + font-weight: 400; + line-height: 1.55; + padding: 8px 10px; + border-radius: 6px; + pointer-events: none; + z-index: 300; + white-space: normal; +} + +.ps-info-tooltip::after { + content: ''; + position: absolute; + top: 50%; + left: 100%; + transform: translateY(-50%); + border: 5px solid transparent; + border-left-color: #111; +} + +.ps-info-wrap:hover .ps-info-tooltip { + display: block; +} + + .pipeline-select { width: 100%; padding: 10px 14px; @@ -240,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; @@ -255,6 +581,7 @@ /* Run button */ .run-btn { flex-shrink: 0; + margin-left: 12px; padding: 10px 32px; background: #002D72; border: 1px solid #002D72; @@ -277,6 +604,26 @@ cursor: not-allowed; } +/* ── Stop button ── */ +.stop-btn { + padding: 7px 18px; + background: rgba(0,0,0,.04); + border: 1px solid rgba(0,0,0,.12); + border-radius: 8px; + color: #6a6a6a; + font-family: 'Space Grotesk', sans-serif; + font-size: 12px; + font-weight: 600; + cursor: pointer; + transition: all 0.15s; + letter-spacing: 0.02em; +} +.stop-btn:hover { + background: rgba(0,0,0,.07); + border-color: rgba(0,0,0,.2); + color: #111; +} + /* ── Advanced options ── */ .advanced-section { margin-top: 20px; @@ -444,7 +791,7 @@ .result-btn:hover { background: rgba(0,0,0,.1); border-color: rgba(0,0,0,.25); - color: #fff; + color: #111111; } .result-btn-primary { background: #002D72; @@ -474,6 +821,208 @@ color: #6a6a6a; } +/* ── Processing counter ── */ +.upload-processing-counter { + display: inline-flex; + align-items: center; + gap: 8px; + font-family: 'JetBrains Mono', monospace; + font-size: 12px; + color: #6a6a6a; + letter-spacing: 0.04em; +} +.upload-processing-dot { + width: 7px; + height: 7px; + border-radius: 50%; + background: #002D72; + animation: upload-pulse 1.6s ease-in-out infinite; + flex-shrink: 0; +} +@keyframes upload-pulse { + 0%, 100% { opacity: 1; transform: scale(1); } + 50% { opacity: 0.4; transform: scale(0.8); } +} + +/* ── Completed rows ── */ +.upload-completed-row { + display: flex; + align-items: center; + gap: 12px; + padding: 12px 14px; + background: #f5f5f5; + border: 1px solid rgba(0,0,0,.06); + border-radius: 10px; + cursor: pointer; + transition: all 0.2s; +} +.upload-completed-row:hover { + background: rgba(0,0,0,.04); + border-color: rgba(0,0,0,.12); +} + +/* ── Download-all button ── */ +.upload-dl-all-btn { + display: inline-flex; + align-items: center; + gap: 6px; + padding: 6px 14px; + background: rgba(0,0,0,.06); + border: 1px solid rgba(0,0,0,.12); + border-radius: 8px; + color: #111111; + font-family: 'Space Grotesk', sans-serif; + font-size: 11px; + font-weight: 600; + cursor: pointer; + transition: all 0.2s; + letter-spacing: 0.02em; +} +.upload-dl-all-btn:hover { + background: rgba(0,0,0,.1); + border-color: rgba(0,0,0,.2); +} + +/* ── File chip internals ── */ +.file-chip-name { + flex: 1; + min-width: 0; + overflow: hidden; + text-overflow: ellipsis; + white-space: nowrap; +} +.file-chip-preview { + padding: 2px 8px; + background: rgba(0,45,114,.1); + border: 1px solid rgba(0,45,114,.2); + border-radius: 10px; + color: #002D72; + font-family: 'Space Grotesk', sans-serif; + font-size: 10px; + font-weight: 600; + cursor: pointer; + letter-spacing: 0.04em; + transition: all 0.15s; + flex-shrink: 0; +} +.file-chip-preview:hover { + background: rgba(0,45,114,.18); +} +.file-chip--active { + border-color: rgba(0,45,114,.3); + background: rgba(0,45,114,.06); +} + +/* ── Pending CT preview panel ── */ +.pending-preview-panel { + margin-bottom: 20px; +} + +/* ── Active job progress card ── */ +.upload-active-job { + background: #f5f5f5; + border: 1px solid rgba(0,0,0,.06); + border-radius: 10px; + padding: 14px 16px; + margin-bottom: 8px; +} +.upload-active-job-header { + display: flex; + align-items: center; + gap: 8px; + margin-bottom: 12px; +} +.upload-active-job-name { + flex: 1; + min-width: 0; + font-family: 'Space Grotesk', sans-serif; + font-size: 13px; + font-weight: 600; + color: #111111; + overflow: hidden; + text-overflow: ellipsis; + white-space: nowrap; +} +.upload-active-job-status { + font-family: 'JetBrains Mono', monospace; + font-size: 10px; + color: #6a6a6a; + letter-spacing: 0.06em; + flex-shrink: 0; +} + +/* ── Progress rows ── */ +.upload-progress-row { + display: flex; + align-items: center; + gap: 10px; + margin-top: 8px; +} +.upload-progress-label { + font-family: 'JetBrains Mono', monospace; + font-size: 10px; + letter-spacing: 0.08em; + text-transform: uppercase; + color: #8f8f8f; + width: 58px; + flex-shrink: 0; +} +.upload-progress-track { + flex: 1; + height: 5px; + background: rgba(0,0,0,.06); + border-radius: 3px; + overflow: hidden; +} +.upload-progress-fill { + height: 100%; + border-radius: 3px; + transition: width 0.35s cubic-bezier(.2,.8,.2,1); +} +.upload-progress-fill--upload { + background: linear-gradient(90deg, #8f8f8f, #002D72); +} +.upload-progress-fill--inference { + background: linear-gradient(90deg, #6a6a6a, #00399a); +} +.upload-progress-pct { + font-family: 'JetBrains Mono', monospace; + font-size: 10px; + font-weight: 600; + color: #111111; + width: 30px; + text-align: right; + flex-shrink: 0; +} + +/* ── Failed row variant ── */ +.upload-completed-row--failed { + cursor: default; + border-color: rgba(0,0,0,.08); + background: rgba(0,0,0,.03); +} +.upload-completed-row--failed:hover { + background: rgba(0,0,0,.05); + border-color: rgba(0,0,0,.12); +} + +/* ── Completed row icon wrapper ── */ +.upload-completed-icon { + width: 36px; + height: 36px; + border-radius: 8px; + background: rgba(0,0,0,.06); + border: 1px solid rgba(0,0,0,.08); + display: flex; + align-items: center; + justify-content: center; + flex-shrink: 0; +} +.upload-completed-row--failed .upload-completed-icon { + background: rgba(0,0,0,.04); + border-color: rgba(0,0,0,.08); +} + /* ── Pre-inference CT preview ── */ .ct-preview-label { font-family: 'Space Grotesk', sans-serif; @@ -520,3 +1069,113 @@ text-align: center; padding: 0 16px; } + +/* ── DICOM slice controls (inline stack preview) ── */ +.dicom-preview-controls { + position: absolute; + left: 12px; + right: 12px; + bottom: 12px; + display: flex; + align-items: center; + gap: 10px; + padding: 6px 10px; + background: rgba(0, 0, 0, 0.45); + border-radius: 8px; + backdrop-filter: blur(2px); +} +.dicom-preview-slider { + flex: 1; + accent-color: #4a90e2; + cursor: pointer; +} +.dicom-preview-count { + font-family: 'JetBrains Mono', monospace; + font-size: 11px; + color: rgba(255, 255, 255, 0.75); + white-space: nowrap; + 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; +} + +/* ── Simplified config: Auto hint, speed/quality toggle, result note ── */ +.simple-config { + display: flex; + flex-direction: column; + gap: 12px; + margin-top: 16px; +} + +.auto-hint-input { + width: 100%; + padding: 11px 14px; + border: 1px solid rgba(0, 45, 114, 0.16); + border-radius: 10px; + font-size: 13.5px; + font-family: inherit; + color: #1a2b47; + background: #fbfcfe; + outline: none; + transition: border-color 0.15s ease, box-shadow 0.15s ease; +} +.auto-hint-input:focus { + border-color: #002D72; + box-shadow: 0 0 0 3px rgba(0, 45, 114, 0.08); +} + +.precision-toggle { + display: grid; + grid-template-columns: 1fr 1fr; + gap: 10px; +} +.precision-opt { + display: flex; + flex-direction: column; + align-items: flex-start; + gap: 3px; + padding: 12px 14px; + border: 1px solid rgba(0, 45, 114, 0.16); + border-radius: 10px; + background: #fff; + cursor: pointer; + text-align: left; + transition: border-color 0.15s ease, background 0.15s ease, box-shadow 0.15s ease; +} +.precision-opt:hover { border-color: rgba(0, 45, 114, 0.4); } +.precision-opt.active { + border-color: #002D72; + background: rgba(0, 45, 114, 0.05); + box-shadow: 0 0 0 2px rgba(0, 45, 114, 0.12); +} +.precision-opt-title { + font-size: 13.5px; + font-weight: 600; + color: #002D72; +} +.precision-opt-sub { + font-size: 11.5px; + color: #6b7793; +} + +.auto-note { + padding: 10px 13px; + border-radius: 10px; + background: rgba(45, 130, 90, 0.08); + border: 1px solid rgba(45, 130, 90, 0.2); + color: #1f5b3f; + font-size: 12.5px; + line-height: 1.45; +} diff --git a/PanTS-Demo/src/routes/UploadPage.tsx b/PanTS-Demo/src/routes/UploadPage.tsx index 1411db9..ca4714a 100644 --- a/PanTS-Demo/src/routes/UploadPage.tsx +++ b/PanTS-Demo/src/routes/UploadPage.tsx @@ -1,19 +1,41 @@ import React, { lazy, Suspense, useCallback, useEffect, useRef, useState } from 'react'; + +const MODEL_OPTIONS: { id: string; label: string; desc: string }[] = [ + { id: "Auto", label: "Auto — AI picks the best model", desc: "Recommended · a local AI reads your request and chooses for you" }, + { id: "None", label: "None", desc: "View only — files never leave your browser" }, + { 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'; -// Lazy so NiiVue isn't pulled into the upload bundle until a file is actually previewed. +// Lazy so NiiVue / Cornerstone aren't pulled into the upload bundle until a file +// is actually previewed. CtPreview handles NIfTI, DicomPreview handles a DICOM series. const CtPreview = lazy(() => import('../components/CtPreview/CtPreview')); +const DicomPreview = lazy(() => import('../components/CtPreview/DicomPreview')); 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 { setLocalNiftiFile } from '../helpers/localNifti'; +import { + deletePendingUpload, + loadPendingUploads, + savePendingUpload, + setPendingNextChunk, + type PendingUpload, +} from '../helpers/pendingUploads'; const parseApiResponse = async (res: Response): Promise => { const contentType = res.headers.get("content-type") || ""; @@ -27,26 +49,67 @@ const parseApiResponse = async (res: Response): Promise => { ); }; +// A selection is either a single NIfTI file or a picked DICOM folder (the series' +// raw .dcm slices). Both are previewable individually and runnable through inference. +type SelectedItem = + | { id: string; kind: 'nifti'; file: File } + | { id: string; kind: 'dicom'; files: File[]; label: string }; + const UploadPage: React.FC = () => { const navigate = useNavigate(); const fileInputRef = useRef(null); - const inferencePollRef = useRef | null>(null); + // Folder picker for a DICOM series (run inference, or view-only when model is "None"). + const dicomUploadInputRef = useRef(null); + // 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); - const [selectedFiles, setSelectedFiles] = useState([]); + const [selectedItems, setSelectedItems] = useState([]); + // Which selected item's inline preview is open (null = none). One at a time. + const [previewItemId, setPreviewItemId] = useState(null); const [message, setMessage] = useState(""); const [serverPath, setServerPath] = useState(""); const [sessionId, setSessionId] = useState(""); - const [uploadedFilename, setUploadedFilename] = useState(""); const [bdmapId, setBdmapId] = useState(""); const [uploadProgress, setUploadProgress] = useState(0); const [isUploading, setIsUploading] = useState(false); - const [inferenceProgress, setInferenceProgress] = useState(0); - const [isInferencing, setIsInferencing] = useState(false); const [inferenceCompleted, setInferenceCompleted] = useState(false); - const [selectedModel, setSelectedModel] = useState<"ePAI" | "SuPreM" | "OpenVAE" | "MedFormer" | "R-Super" | "Atlas-Net" | "">(""); + const [selectedModel, setSelectedModel] = useState<"Auto" | "None" | "ePAI" | "SuPreM" | "OpenVAE" | "MedFormer" | "R-Super" | "Atlas-Net" | "">("Auto"); + // FP16 = quantized (faster, half the memory, identical masks); the default. + const [selectedPrecision, setSelectedPrecision] = useState<"fp16" | "fp32">("fp16"); + // Optional plain-language hint for the Auto picker, and the note it returns. + const [autoHint, setAutoHint] = useState(""); + const [autoNote, setAutoNote] = useState(""); + const [autoBusy, setAutoBusy] = useState(false); + 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"]; @@ -60,7 +123,7 @@ const UploadPage: React.FC = () => { alert("Please select .nii or .nii.gz files only"); return; } - setSelectedFiles(prev => [...prev, ...filteredFiles]); + setSelectedItems(prev => [...prev, ...filteredFiles.map(f => ({ id: crypto.randomUUID(), kind: 'nifti' as const, file: f }))]); }; const handleDrop = useCallback((e: React.DragEvent) => { @@ -74,7 +137,7 @@ const UploadPage: React.FC = () => { alert("Please drop .nii or .nii.gz files only"); return; } - setSelectedFiles(prev => [...prev, ...filteredFiles]); + setSelectedItems(prev => [...prev, ...filteredFiles.map(f => ({ id: crypto.randomUUID(), kind: 'nifti' as const, file: f }))]); }, []); const handleDragOver = useCallback((e: React.DragEvent) => { @@ -86,182 +149,501 @@ const UploadPage: React.FC = () => { setIsDragOver(false); }, []); - const removeFile = (index: number) => { - setSelectedFiles(prev => prev.filter((_, i) => i !== index)); + const removeItem = (id: string) => { + setSelectedItems(prev => prev.filter(item => item.id !== id)); + setPreviewItemId(prev => (prev === id ? null : prev)); }; - /* ── Inference polling ── */ - const stopInferencePolling = () => { - if (inferencePollRef.current) { - clearInterval(inferencePollRef.current); - inferencePollRef.current = null; + // Pick a DICOM folder to RUN INFERENCE on (distinct from the view-only opener): + // the raw slices are added as one selectable item, previewable and runnable. + const handleDicomInferenceSelect = (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; } + setSelectedItems(prev => [ + ...prev, + { id: crypto.randomUUID(), kind: 'dicom', files: candidates, label: `DICOM series (${candidates.length} slices)` }, + ]); }; - const startInferencePolling = (sid: string, model: string) => { - stopInferencePolling(); - setIsInferencing(true); - setInferenceProgress(5); + /* ── 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); + } + }; - inferencePollRef.current = setInterval(async () => { + 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); - if (!res.ok) throw new Error(data.error || data.status || "Status check failed"); - const status = (data.status || "").toLowerCase(); - if (status === "completed") { - setInferenceProgress(100); - setIsInferencing(false); - setInferenceCompleted(true); - setRecentUploads(updateRecentUploadStatus(sid, "Completed")); - stopInferencePolling(); - - setTimeout(() => { - if (model === "OpenVAE") { - navigate(`/reconstruction/${sid}`); - } else { - navigate(`/session/${sid}`); - } - }, 600); + + // 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; } - if (status === "failed") { - setIsInferencing(false); + 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")); - stopInferencePolling(); setMessage(`Inference failed${data.error ? `: ${data.error}` : ""}`); - return; + } 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); } - setInferenceProgress(prev => Math.min(95, Math.max(prev + 7, 10))); } catch (err) { - setInferenceProgress(prev => Math.min(95, Math.max(prev + 3, 10))); + // 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}`); }; useEffect(() => { - return () => { stopInferencePolling(); }; + // 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(); }; }, []); - /* ── Upload (chunked) ── */ - const CHUNK_SIZE = 256 * 1024; + // 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]); - /* ── Run inference ── */ - const handleRunEpaiInference = async () => { - if (!sessionId && !serverPath.trim() && selectedFiles.length === 0) { - alert("Provide a server file path or upload/select a file first."); - return; - } + 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]); - let currentSessionId = sessionId; - let currentUploadedFilename = uploadedFilename; + 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]); - try { - // If files selected but not yet uploaded, upload first - if (!currentSessionId && !serverPath.trim() && selectedFiles.length > 0) { - const file = selectedFiles[0]; - const totalChunks = Math.ceil(file.size / CHUNK_SIZE); - const newSessionId = crypto.randomUUID(); + 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]); - setIsUploading(true); - setUploadProgress(0); - setMessage(`Uploading ${file.name}...`); - - for (let i = 0; i < totalChunks; i++) { - const start = i * CHUNK_SIZE; - const end = Math.min(start + CHUNK_SIZE, file.size); - const chunk = file.slice(start, end); - - const formData = new FormData(); - formData.append("session_id", newSessionId); - 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, - }); - - if (res.status === 413) { - throw new Error("Upload chunk too large for server/proxy limit (HTTP 413)."); - } + /* ── Upload (chunked) ── */ + const CHUNK_SIZE = 256 * 1024; - const data = await parseApiResponse(res); - if (!res.ok) throw new Error(data.error || "Chunk upload failed"); - setUploadProgress(Math.round(((i + 1) / totalChunks) * 100)); - } + // 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 { + if (foreground) { + foregroundUploadSidRef.current = sid; + setIsUploading(true); + setUploadProgress(Math.round((p.nextChunk / totalChunks) * 100)); + setMessage(`Uploading ${filename}...`); + } - setMessage("Finalizing upload..."); + for (let i = p.nextChunk; i < totalChunks; i++) { + const chunk = file.slice(i * CHUNK_SIZE, Math.min((i + 1) * CHUNK_SIZE, file.size)); + 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 finalizeRes = await fetch(`${API_BASE}/api/finalize-upload`, { - method: "POST", - body: new URLSearchParams({ - session_id: newSessionId, - total_chunks: totalChunks.toString(), - output_filename: file.name, - ...(bdmapId.trim() ? { bdmap_id: bdmapId.trim() } : {}), - }), + 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"); - const finalizeData = await parseApiResponse(finalizeRes); - if (!finalizeRes.ok) throw new Error(finalizeData.error); + // 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", + signal: controller.signal, + body: new URLSearchParams({ + session_id: sid, + total_chunks: totalChunks.toString(), + output_filename: filename, + ...(bid ? { bdmap_id: bid } : {}), + }), + }); + const finalizeData = await parseApiResponse(finalizeRes); + if (!finalizeRes.ok) throw new Error(finalizeData.error); + const uploadedName = finalizeData.uploaded_filename || filename; - currentSessionId = newSessionId; - currentUploadedFilename = finalizeData.uploaded_filename || file.name; - - setSessionId(currentSessionId); - setUploadedFilename(currentUploadedFilename); - setServerPath(finalizeData.path || ""); + // 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); } - setMessage(`Starting ${selectedModel} inference...`); - setInferenceProgress(0); - setIsInferencing(true); + if (foreground) setMessage(`Starting ${model} inference...`); + const inferFd = new FormData(); + inferFd.append("session_id", sid); + inferFd.append("model_name", model); + inferFd.append("precision", selectedPrecision); + 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) { + // 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); + } + } + }; - const formData = new FormData(); - formData.append("session_id", currentSessionId || crypto.randomUUID()); - formData.append("model_name", selectedModel); + // DICOM folder → inference. Uploads each raw slice, asks the server to convert + // the series to NIfTI (SimpleITK), then hands off to the same inference + polling + // flow as a NIfTI run. Not IndexedDB-resumable (a folder is many files); a reload + // mid-upload marks the run Failed, consistent with the Active/Recent cards. Wired + // into uploadAbortRef so the Active card's Cancel button aborts it cleanly. + const runDicomUpload = async (sid: string, files: File[], model: string) => { + const controller = new AbortController(); + uploadAbortRef.current.set(sid, controller); + foregroundUploadSidRef.current = sid; + uploadResumableRef.current = false; // no IDB copy for a DICOM folder + setPhase(sid, "uploading"); + try { + setIsUploading(true); + setUploadProgress(0); + setMessage(`Uploading DICOM series (${files.length} slices)...`); - if (serverPath.trim()) { - formData.append("INPUT_SERVER_PATH", serverPath.trim()); - } else if (currentUploadedFilename) { - formData.append("uploaded_filename", currentUploadedFilename); + for (let i = 0; i < files.length; i++) { + const formData = new FormData(); + formData.append("session_id", sid); + formData.append("file", files[i]); + const res = await fetch(`${API_BASE}/api/upload-dicom-slice`, { + method: "POST", body: formData, signal: controller.signal, + }); + if (res.status === 413) throw new Error("DICOM slice too large for server/proxy limit (HTTP 413)."); + const data = await parseApiResponse(res); + if (!res.ok) throw new Error(data.error || "DICOM slice upload failed"); + setUploadProgress(Math.round(((i + 1) / files.length) * 100)); } + setMessage("Converting DICOM series to NIfTI..."); + const finalizeRes = await fetch(`${API_BASE}/api/finalize-dicom`, { + method: "POST", signal: controller.signal, + body: new URLSearchParams({ session_id: sid }), + }); + const finalizeData = await parseApiResponse(finalizeRes); + if (!finalizeRes.ok) throw new Error(finalizeData.error || "DICOM conversion failed"); + const uploadedName = finalizeData.uploaded_filename || "ct.nii.gz"; + + foregroundUploadSidRef.current = null; + setUploadProgress(100); + setIsUploading(false); + + setMessage(`Starting ${model} inference...`); + const inferFd = new FormData(); + inferFd.append("session_id", sid); + inferFd.append("model_name", model); + inferFd.append("precision", selectedPrecision); + inferFd.append("uploaded_filename", uploadedName); const res = await fetch(`${API_BASE}/api/run-epai-inference`, { - method: "POST", - body: formData, + method: "POST", body: inferFd, signal: controller.signal, }); const data = await parseApiResponse(res); if (!res.ok) throw new Error(data.error || "Failed to start inference"); - const sid = data.session_id || formData.get("session_id")?.toString() || ""; setSessionId(sid); - setMessage(`${selectedModel} inference started. Session: ${sid}`); - if (sid) { - setRecentUploads( - addRecentUpload({ - sessionId: sid, - label: bdmapId.trim() || currentUploadedFilename || selectedFiles[0]?.name || sid, - model: selectedModel, - status: "Processing", - timestamp: Date.now(), - isReconstruction: selectedModel === "OpenVAE", - }) - ); - startInferencePolling(sid, selectedModel); - } + setPhase(sid, "queued"); // server queues for the GPU; poll refines this + setMessage(`${model} inference started. Session: ${sid}`); + startInferencePolling(sid, model); } catch (err) { + // A user cancel aborts our fetches - cancelRun already set the card to + // Cancelled, so don't overwrite that with Failed. + if (controller.signal.aborted) return; console.error(err); + setPhase(sid); + foregroundUploadSidRef.current = null; setIsUploading(false); - setIsInferencing(false); + setRecentUploads(updateRecentUploadStatus(sid, "Failed")); setMessage("Failed: " + (err as Error).message); + } finally { + if (uploadAbortRef.current.get(sid) === controller) { + uploadAbortRef.current.delete(sid); + } + } + }; + + /* ── Run inference ── */ + const handleRunEpaiInference = async () => { + const item = selectedItems[0] ?? null; + + // "None" model = view only: open the scan in its full local viewer, nothing is + // uploaded or run. DICOM opens the /dicom viewer, NIfTI the /local-nifti viewer. + if (selectedModel === "None") { + if (!item) { alert("Select a scan to view first."); return; } + if (item.kind === "dicom") { + setLocalDicomFiles(item.files); + navigate("/dicom"); + } else { + setLocalNiftiFile(item.file); + navigate("/local-nifti"); + } + return; + } + + const path = serverPath.trim(); + if (!item && !path) { + alert("Provide a server file path or upload/select a file first."); + return; + } + + // Auto mode: ask the backend (local LLM, with a heuristic fallback) to choose + // the model + precision, then run with that choice. Resolves to a concrete + // model name before anything is uploaded. + let model: string = selectedModel; + let precision = selectedPrecision; + if (selectedModel === "Auto") { + try { + setAutoBusy(true); + setMessage("Auto-selecting the best model for your scan…"); + const filename = item ? (item.kind === 'dicom' ? item.label : item.file.name) : (path ? path.split("/").pop() : ""); + const res = await fetch(`${API_BASE}/api/suggest-model`, { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify({ hint: autoHint, filename }), + }); + const data = await parseApiResponse(res); + model = data.model || "ePAI"; + precision = (data.precision === "fp32" ? "fp32" : "fp16"); + setSelectedModel(model as typeof selectedModel); + setSelectedPrecision(precision); + setAutoNote(`Auto-selected ${model} · ${precision.toUpperCase()} — ${data.reason || "best fit for this scan."}`); + } catch { + model = "ePAI"; + precision = "fp16"; + setSelectedModel("ePAI"); + setAutoNote("Auto-select was unavailable, so it defaulted to ePAI · FP16."); + } finally { + setAutoBusy(false); + } + } + + const sid = crypto.randomUUID(); + const itemName = item ? (item.kind === 'dicom' ? item.label : item.file.name) : undefined; + const label = bdmapId.trim() || (path ? path.split("/").pop() : itemName) || 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("model_name", model); + fd.append("precision", selectedPrecision); + 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 || "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; } + + // Consume the first item so the next can be queued. A DICOM folder uploads its + // slices and converts server-side; a NIfTI file rides the resumable path (stashed + // in IndexedDB so an interrupted upload can resume). + setSelectedItems(prev => prev.slice(1)); + + if (item!.kind === 'dicom') { + runDicomUpload(sid, item!.files, model); + return; + } + + const file = item!.file; + const pending: PendingUpload = { + sessionId: sid, + 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 handleCheckStatus = async () => { @@ -273,13 +655,20 @@ const UploadPage: React.FC = () => { setMessage(`Status: ${data.status}${data.error ? ` (${data.error})` : ""}`); const status = (data.status || "").toLowerCase(); if (status === "completed") { - setInferenceProgress(100); - setIsInferencing(false); setInferenceCompleted(true); setRecentUploads(updateRecentUploadStatus(sessionId, "Completed")); - stopInferencePolling(); - } else if (status === "running") { - if (!isInferencing) startInferencePolling(sessionId, selectedModel); + 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) { console.error(err); @@ -298,9 +687,7 @@ const UploadPage: React.FC = () => { setMessage(`Status: ${statusData.status || "unknown"}. Please wait until completed.`); return; } - setInferenceProgress(100); - setIsInferencing(false); - stopInferencePolling(); + stopPolling(sessionId); const resultRes = await fetch(`${API_BASE}/api/get_result/${sessionId}`); if (!resultRes.ok) { @@ -330,12 +717,12 @@ const UploadPage: React.FC = () => { } const newSessionId = crypto.randomUUID(); setInferenceCompleted(false); - setInferenceProgress(0); setMessage("Starting ePAI inference on reconstructed CT..."); const formData = new FormData(); formData.append("session_id", newSessionId); formData.append("model_name", "ePAI"); + formData.append("precision", selectedPrecision); formData.append("source_reconstruction_session_id", sessionId); try { @@ -350,7 +737,18 @@ const UploadPage: React.FC = () => { setSessionId(sid); setSelectedModel("ePAI" as const); setMessage(`ePAI inference started on reconstructed CT. Session: ${sid}`); - if (sid) startInferencePolling(sid, "ePAI"); + if (sid) { + setRecentUploads( + addRecentUpload({ + sessionId: sid, + label: "ePAI on reconstruction", + model: "ePAI", + status: "Processing", + timestamp: Date.now(), + }) + ); + startInferencePolling(sid, "ePAI"); + } } catch (err) { console.error(err); setMessage("Failed to start ePAI on reconstruction: " + (err as Error).message); @@ -358,6 +756,8 @@ const UploadPage: React.FC = () => { }; /* ── Render ── */ + const previewItem = selectedItems.find(i => i.id === previewItemId) ?? null; + return (
{/* Ambient glow */} @@ -388,33 +788,78 @@ const UploadPage: React.FC = () => { style={{ display: 'none' }} onChange={handleFileSelect} /> + { + dicomUploadInputRef.current = el; + if (el) { + el.setAttribute('webkitdirectory', ''); + el.setAttribute('directory', ''); + } + }} + type="file" + multiple + style={{ display: 'none' }} + onChange={handleDicomInferenceSelect} + />
Click or drag to upload
-
.nii or .nii.gz
+
+ + +
- {/* ── File chips ── */} - {selectedFiles.length > 0 && ( + {/* ── Selected items: NIfTI files + DICOM series, each individually previewable ── */} + {selectedItems.length > 0 && (
- {selectedFiles.map((file, index) => ( -
- {file.name} - -
- ))} + {selectedItems.map((item) => { + const name = item.kind === 'dicom' ? item.label : item.file.name; + const isOpen = previewItemId === item.id; + return ( +
+ {name} + + +
+ ); + })}
)} {/* ── Pre-inference preview: inspect the selected scan before running a model ── */} - {selectedFiles.length > 0 && !isUploading && !isInferencing && !inferenceCompleted && ( + {previewItem && !isUploading && ( <> -
Preview · {selectedFiles[0].name}
+
+ Preview · {previewItem.kind === 'dicom' ? previewItem.label : previewItem.file.name} +
Loading preview…}> - + {previewItem.kind === 'dicom' + ? + : } )} @@ -428,10 +873,44 @@ const UploadPage: React.FC = () => { 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 && ( + + + + )} +
+
+ ))} +
+ )} +
→
@@ -442,18 +921,41 @@ const UploadPage: React.FC = () => {
2
Model - +
+ + {modelDropOpen && ( +
+ {MODEL_OPTIONS.map(m => ( +
{ setSelectedModel(m.id as typeof selectedModel); setModelDropOpen(false); }} + > +
+ {m.label} + {m.desc} +
+
+ {selectedModel === m.id && ( + + + + )} +
+
+ ))} +
+ )} +
→
@@ -465,21 +967,87 @@ const UploadPage: React.FC = () => { 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 && ( + + + + )} +
+
+ ))} +
+ )} +
+ {/* ── Simple controls: Auto hint, speed/quality, result note ── */} +
+ {selectedModel === "Auto" && ( + setAutoHint(e.target.value)} + /> + )} +
+ + +
+ {autoNote &&
{autoNote}
} +
+ {/* ── Advanced options ── */}
)} - {/* ── Progress ── */} - {(isUploading || uploadProgress > 0) && ( + {/* ── Progress (upload phase only - running inference shows in Active below) ── */} + {isUploading && (
@@ -534,20 +1102,6 @@ const UploadPage: React.FC = () => {
)} - {(isInferencing || inferenceProgress > 0) && ( -
-
-
- Inference Progress - {inferenceProgress}% -
-
-
-
-
-
- )} - {/* ── Results ── */} {inferenceCompleted && sessionId && (
@@ -586,6 +1140,70 @@ const UploadPage: React.FC = () => { {message &&
{message}
}
+ {/* ── 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} + + +
+
+ ); + })} +
+
+ )} + {/* ── Recent Uploads ── */}
{ Recent Uploads
- {recentUploads.length === 0 ? ( + {recentUploads.filter(u => u.status !== "Processing").length === 0 ? (
{ 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.map((upload) => { + recentUploads.filter(u => u.status !== "Processing").map((upload) => { + const clickable = upload.status !== "Failed" && upload.status !== "Cancelled"; const openSession = () => { - if (upload.status === "Failed") return; + if (!clickable) return; navigate(`/${upload.isReconstruction ? "reconstruction" : "session"}/${upload.sessionId}`); }; - const clickable = upload.status !== "Failed"; return (
{ View )} +
); diff --git a/flask-server/api/api_blueprint.py b/flask-server/api/api_blueprint.py index 4edcdb8..0446860 100644 --- a/flask-server/api/api_blueprint.py +++ b/flask-server/api/api_blueprint.py @@ -1,14 +1,24 @@ from flask import Blueprint, send_file, make_response, request, jsonify, Response +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 +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 +from services.ollama_client import ( + DEFAULT_OLLAMA_MODEL, + OllamaUnavailable, + chat_json, + list_ollama_models, +) +from services.segmentation_metrics import calculate_session_metrics from models.application_session import ApplicationSession from models.combined_labels import CombinedLabels from models.base import db from constants import Constants import zipfile +import json import pandas as pd from pathlib import Path @@ -35,6 +45,25 @@ api_blueprint = Blueprint("api", __name__) last_session_check = datetime.now() +# Low-res volumes (generated by scripts/make_lowres.py) live on a WRITABLE disk, +# NOT the read-only dataset mount. The tree mirrors the dataset: +# /image_only//ct_lowres.nii.gz +# /mask_only//combined_labels_lowres.nii.gz +# Overridable via PANTS_LOWRES_PATH. If a file is absent the API serves full res, +# so this stays fully additive (no-op until the batch has run). +LOWRES_ROOT = os.environ.get("PANTS_LOWRES_PATH", "/home/visitor/pants_lowres") + +import hmac +import threading + +# Session/case ids come straight from client requests and are joined into +# filesystem paths below. _is_safe_id (pure, unit-tested in tests/unit/ +# test_path_safety.py) rejects anything that could escape the intended +# directory before it touches os.path; secure_filename is the barrier at each +# path-construction site. +from .path_safety import is_safe_id as _is_safe_id + + def _load_metadata_cache(): try: xlsx_path = os.path.join(Constants.PANTS_PATH, "metadata.xlsx") @@ -81,7 +110,7 @@ def _require_worker_auth(): return jsonify({"error": "WORKER_API_TOKEN is not configured on server"}), 500 provided = (request.headers.get("X-Worker-Token", "") or "").strip() - if provided != expected: + if not hmac.compare_digest(provided, expected): return jsonify({"error": "Unauthorized worker token"}), 401 return None @@ -179,7 +208,9 @@ def get_preview(clabel_ids): # if not preloaded @api_blueprint.route('/get_image_preview/', methods=['GET']) def get_image_preview(clabel_id): - path = os.path.join(Constants.PANTS_PATH, "profile_only", get_panTS_id(clabel_id), "profile.jpg") + if not _is_safe_id(clabel_id): + return jsonify({"error": "Invalid id"}), 400 + path = os.path.join(Constants.PANTS_PATH, "profile_only", get_panTS_id(secure_filename(clabel_id)), "profile.jpg") if not os.path.exists(path): return jsonify({"error": f"File not found: {path} "}), 404 return send_file( @@ -192,7 +223,9 @@ def get_image_preview(clabel_id): @api_blueprint.route("/cases//mesh-manifest") def get_mesh_manifest(case_id): - manifest_path = os.path.join(Constants.MESH_PATH, get_panTS_id(case_id), "manifest.json") + if not _is_safe_id(case_id): + return jsonify({"error": "Invalid id"}), 400 + manifest_path = os.path.join(Constants.MESH_PATH, get_panTS_id(secure_filename(case_id)), "manifest.json") if not os.path.exists(manifest_path): return jsonify({"error": f"File not found: {manifest_path} "}), 404 @@ -270,7 +303,9 @@ def upload(): session_id = request.form.get('SESSION_ID') if not session_id: return jsonify({"error": "No session ID provided"}), 400 - + if not _is_safe_id(session_id): + return jsonify({"error": "Invalid session ID"}), 400 + base_path = os.path.join(Constants.SESSIONS_DIR_NAME, session_id) os.makedirs(base_path, exist_ok=True) @@ -309,15 +344,19 @@ def get_mask_data(): return jsonify(result) -@api_blueprint.route('/get-main-nifti/', methods=['GET']) +@api_blueprint.route('/get-main-nifti/.nii.gz', methods=['GET']) def get_main_nifti(clabel_id): - case_dir = f"{Constants.PANTS_PATH}/image_only/{get_panTS_id(clabel_id)}" + if not _is_safe_id(clabel_id): + return jsonify({"error": "Invalid id"}), 400 + case_dir = f"{Constants.PANTS_PATH}/image_only/{get_panTS_id(secure_filename(clabel_id))}" main_nifti_path = f"{case_dir}/{Constants.MAIN_NIFTI_FILENAME}" # ?res=low → serve the precomputed low-res copy when present (much smaller/faster - # for big full-body scans). Falls back to full res if it hasn't been generated. + # for big full-body scans). It lives under LOWRES_ROOT (a writable disk), NOT the + # read-only dataset mount. Falls back to full res if it hasn't been generated. if (request.args.get('res') or '').strip().lower() == 'low': - low_path = f"{case_dir}/{Constants.MAIN_NIFTI_FILENAME.replace('.nii.gz', '_lowres.nii.gz')}" + low_name = Constants.MAIN_NIFTI_FILENAME.replace('.nii.gz', '_lowres.nii.gz') + low_path = f"{LOWRES_ROOT}/image_only/{get_panTS_id(secure_filename(clabel_id))}/{low_name}" if os.path.exists(low_path): main_nifti_path = low_path @@ -326,7 +365,6 @@ def get_main_nifti(clabel_id): response.headers['Cross-Origin-Resource-Policy'] = 'cross-origin' # response.headers['Cross-Origin-Embedder-Policy'] = 'require-corp' - response.headers['Content-Encoding'] = 'gzip' # Volumes are immutable per case — let the browser cache so revisits are instant. response.headers['Cache-Control'] = 'public, max-age=604800, immutable' @@ -349,14 +387,17 @@ def get_report(id): organ_metrics = organ_metrics.get("organ_metrics", []) except Exception as e: return jsonify({"error": f"Error loading organ metrics: {str(e)}"}), 500 - - + base_path = f"{SESSIONS_DIR}/{id}" + # New flat structure, matching get-main-nifti / get-label-colormap above — + # this fixes a bug from the merge where `subfolder`/`label_subfolder` were + # referenced here but never defined, which would have raised a NameError + # on every call to this route. ct_path = f"{Constants.PANTS_PATH}/image_only/{get_panTS_id(id)}/{Constants.MAIN_NIFTI_FILENAME}" masks = f"{Constants.PANTS_PATH}/mask_only/{get_panTS_id(id)}/{Constants.COMBINED_LABELS_NIFTI_FILENAME}" - + template_pdf = os.getenv("TEMPLATE_PATH", "report_template_3.pdf") - + extracted_data = None column_headers = None try: @@ -366,7 +407,7 @@ def get_report(id): column_headers = df.columns.tolist() except Exception: pass - + generate_pdf_with_template( output_pdf=output_pdf_path, folder_name=id, @@ -378,21 +419,308 @@ def get_report(id): extracted_data=extracted_data, column_headers=column_headers ) - + return send_file( output_pdf_path, mimetype="application/pdf", as_attachment=True, download_name=f"report_{id}.pdf" ) - + except Exception as e: return jsonify({"error": f"Unhandled error: {str(e)}"}), 500 - + finally: if os.path.exists(temp_pdf_path): os.remove(temp_pdf_path) + + +@api_blueprint.route('/explain-impressions', methods=['POST']) +def explain_impressions(): + """Deprecated: plain-language explanations required Ollama, which isn't + available in production. The frontend no longer calls this button, but + the route stays as a stable no-op so old clients don't 404.""" + return jsonify({ + "plain_language": ["Plain-language summary unavailable — please ask your doctor to walk through these findings with you."], + }), 200 + +@api_blueprint.route('/define-term', methods=['GET']) +def define_term(): + """Returns a plain-English definition for a clicked medical term, from + the hardcoded MEDICAL_TERM_DICTIONARY only — no LLM fallback.""" + term = (request.args.get('term') or '').strip().lower() + if not term: + return jsonify({"error": "Missing term parameter"}), 400 + + if term in MEDICAL_TERM_DICTIONARY: + return jsonify({"term": term, "definition": MEDICAL_TERM_DICTIONARY[term], "source": "dictionary"}) + return jsonify({ + "term": term, + "definition": "Definition unavailable — try asking your doctor what this term means.", + "source": "fallback", + }), 200 + + +@api_blueprint.route('/get-report-data/', methods=['GET']) +def get_report_data(id): + if id is None or not str(id).isdigit(): + return jsonify({"error": "Invalid id parameter"}), 400 + case_id = int(id) + try: + if id is None or not str(id).isdigit(): + return jsonify({"error": "Invalid id parameter"}), 400 + id = str(int(id)) + # ── Try RadGPT structured report from metadata.xlsx first ───────────── + # This uses Zongwei Zhou's own RadGPT model output — more accurate + # than Ollama-generated impressions. Falls back to Ollama if not found. + radgpt_comments = None + radgpt_impression = None + try: + import openpyxl, re + metadata_path = os.path.join(Constants.PANTS_PATH, "data", "metadata.xlsx") + if os.path.exists(metadata_path): + wb = openpyxl.load_workbook(metadata_path, read_only=True, data_only=True) + ws = wb.active + headers = [cell.value for cell in next(ws.iter_rows())] + id_col = next((i for i, h in enumerate(headers) if h and 'ID' in str(h)), 0) + report_col = next((i for i, h in enumerate(headers) if h and 'report' in str(h).lower()), -1) + if report_col >= 0: + pants_id = get_panTS_id(id) + for row in ws.iter_rows(min_row=2, values_only=True): + row_id = str(row[id_col] or '').strip() + if row_id == pants_id: + raw = str(row[report_col] or '') + # Clean Windows carriage return artifacts + raw = raw.replace('_x000D_', '\n').replace('\r\n', '\n').replace('\r', '\n') + # Collapse multiple blank lines + import re as _re + raw = _re.sub(r'\n{3,}', '\n\n', raw) + findings_match = re.search(r'FINDINGS:(.*?)(?=IMPRESSION:|$)', raw, re.DOTALL) + impression_match = re.search(r'IMPRESSION:(.*?)$', raw, re.DOTALL) + if findings_match: + radgpt_comments = findings_match.group(1).strip() + if impression_match: + imp_text = impression_match.group(1).strip() + # Keep full impression, split into sentences + sentences = [s.strip() for s in re.split(r'(?<=[.!?])\s+', imp_text) if s.strip()] + radgpt_impression = sentences if sentences else [imp_text] + print(f"[RadGPT] Found report for case {id}: {radgpt_impression}") + break + wb.close() + except Exception as e: + print(f"[RadGPT] metadata lookup failed: {e}") + # ───────────────────────────────────────────────────────────────────── + subfolder = "ImageTr" if case_id < 9000 else "ImageTe" + label_subfolder = "LabelTr" if case_id < 9000 else "LabelTe" + # Check image_only first (new structure), fall back to data/ImageTr + image_only_path = f"{Constants.PANTS_PATH}/image_only/{get_panTS_id(case_id)}/{Constants.MAIN_NIFTI_FILENAME}" + data_ct_path = f"{Constants.PANTS_PATH}/data/{subfolder}/{get_panTS_id(case_id)}/{Constants.MAIN_NIFTI_FILENAME}" + ct_path = image_only_path if os.path.exists(image_only_path) else data_ct_path + # Check mask_only first (new structure), fall back to data/LabelTe + mask_only_path = f"{Constants.PANTS_PATH}/mask_only/{get_panTS_id(case_id)}/{Constants.COMBINED_LABELS_NIFTI_FILENAME}" + data_mask_path = f"{Constants.PANTS_PATH}/data/{label_subfolder}/{get_panTS_id(case_id)}/{Constants.COMBINED_LABELS_NIFTI_FILENAME}" + mask_path = mask_only_path if os.path.exists(mask_only_path) else data_mask_path + seg_dir = f"{Constants.PANTS_PATH}/data/{label_subfolder}/{get_panTS_id(case_id)}/segmentations" + + pid = get_panTS_id(case_id) + meta = _METADATA_CACHE.get(pid, {}) + age = meta.get("age", "N/A") + sex = meta.get("sex", "N/A") + + wb = load_workbook(os.path.join(Constants.PANTS_PATH, "data", "metadata.xlsx")) + sheet = wb["PanTS_metadata"] + contrast = "" + study_detail = "" + for row in sheet.iter_rows(values_only=True): + if row[0] == pid: + contrast = row[3] + study_detail = row[8] + break + + # If local files don't exist, download from HuggingFace + if not os.path.exists(ct_path) or not os.path.exists(mask_path): + import requests, tempfile + pants_id = get_panTS_id(id) + hf_base = f"https://huggingface.co/datasets/BodyMaps/iPanTSMini/resolve/main" + tmp_dir = tempfile.mkdtemp() + if not os.path.exists(ct_path): + hf_ct = f"{hf_base}/image_only/{pants_id}/ct.nii.gz?download=true" + ct_path = os.path.join(tmp_dir, "ct.nii.gz") + print(f"[HuggingFace] Downloading CT for case {id}...") + r = requests.get(hf_ct, stream=True, verify=False) + with open(ct_path, 'wb') as f: + for chunk in r.iter_content(chunk_size=8192): + f.write(chunk) + if not os.path.exists(mask_path): + hf_mask = f"{hf_base}/mask_only/{pants_id}/combined_labels.nii.gz?download=true" + mask_path = os.path.join(tmp_dir, "combined_labels.nii.gz") + print(f"[HuggingFace] Downloading mask for case {id}...") + r = requests.get(hf_mask, stream=True, verify=False) + with open(mask_path, 'wb') as f: + for chunk in r.iter_content(chunk_size=8192): + f.write(chunk) + + ct_nii = nib.load(ct_path) + spacing = ct_nii.header.get_zooms() + shape = ct_nii.shape + ct_array = ct_nii.get_fdata() + mask_nii = nib.load(mask_path) + mask_array = mask_nii.get_fdata().astype(np.uint8) + # Crop both arrays to the minimum shape along each axis + # to handle slight size mismatches between CT and mask + min_shape = tuple(min(c, m) for c, m in zip(ct_array.shape, mask_array.shape)) + ct_array = ct_array[:min_shape[0], :min_shape[1], :min_shape[2]] + mask_array = mask_array[:min_shape[0], :min_shape[1], :min_shape[2]] + voxel_volume = np.prod(mask_nii.header.get_zooms()) / 1000 + # World-space affine - converts a voxel index (i, j, k) to real + # millimeter coordinates. This is what makes the centroid below a + # REAL position usable by moveCornerstoneCrosshairToMm, rather + # than a placeholder. + affine = mask_nii.affine + + LABELS = {v: k for k, v in Constants.PREDEFINED_LABELS.items()} + + # Soft physiological sanity ranges per organ type, used only to + # flag "this needs review" vs "looks normal" internally. NOT + # shown to the user as raw numbers - the frontend only ever sees + # the `status` field. This is a stopgap for a known upstream + # segmentation/data issue where some organs read in the air range; + # rather than silently showing a wrong number as fact, or trying + # to "correct" the data here, we flag it so the UI can say + # "needs review" instead of presenting a confident wrong reading. + SOLID_ORGAN_HU_RANGE = (-20, 150) # liver, spleen, kidney, pancreas, etc. + GI_HOLLOW_ORGAN_HU_RANGE = (-300, 200) # colon, stomach, intestine, duodenum - tightened from + # -1000 so a mean this close to pure air (e.g. -756 for + # colon) correctly flags "check" instead of "normal" - + # a real colon has enough wall/stool tissue that a mean + # in the deep-air range usually signals a segmentation + # issue, not a genuinely normal reading. + LUNG_HU_RANGE = (-1000, -200) # lungs are genuinely air-filled - this range is correct as-is + GI_HOLLOW_ORGANS = {"colon", "stomach", "intestine", "duodenum"} + LUNG_ORGANS = {"lung_left", "lung_right"} + + organ_volumes = {} + NO_FLAG_ORGANS = { + "femur_left", "femur_right", "aorta", "postcava", "veins", + "celiac_artery", "superior_mesenteric_artery", "renal_vein_left", + "renal_vein_right", "common_bile_duct", "pancreatic_duct", + } + for organ, label_id in LABELS.items(): + if label_id == 0: + continue + mask = (mask_array == label_id) + if not np.any(mask): + continue + volume = float(np.sum(mask) * voxel_volume) + mean_hu = float(np.mean(ct_array[mask])) + + if organ in LUNG_ORGANS: + lo, hi = LUNG_HU_RANGE + elif organ in GI_HOLLOW_ORGANS: + lo, hi = GI_HOLLOW_ORGAN_HU_RANGE + else: + lo, hi = SOLID_ORGAN_HU_RANGE + status = "normal" if organ in NO_FLAG_ORGANS else ("check" if (mean_hu < lo or mean_hu > hi) else "normal") + + # Real centroid: voxel-space center of mass, converted to mm + # via the affine. This is genuine anatomical position - not a + # placeholder - and feeds the same crosshair-navigation + # plumbing already used elsewhere (moveCornerstoneCrosshairToMm + # / moveNiiVueCrosshairToMm) for click-to-jump. + voxel_coords = np.argwhere(mask) + centroid_voxel = voxel_coords.mean(axis=0) # (i, j, k) in voxel space + centroid_world = nib.affines.apply_affine(affine, centroid_voxel) + + # Bounding box dimensions in cm (real physical size of the organ) + bbox_min = voxel_coords.min(axis=0) + bbox_max = voxel_coords.max(axis=0) + bbox_voxels = bbox_max - bbox_min + 1 + # Convert voxel counts to mm using spacing, then to cm + spacing_mm = np.abs([affine[0,0], affine[1,1], affine[2,2]]) + dims_mm = bbox_voxels * spacing_mm + dims_cm = [round(float(d)/10, 1) for d in dims_mm] + + organ_volumes[organ] = { + "volume": round(volume, 2), + "mean_hu": round(mean_hu, 1), + "status": status, + "centroid_mm": [round(float(c), 2) for c in centroid_world], + "dimensions": dims_cm, + } + + lesions = {} + lesion_files = { + "pancreas": "pancreatic_lesion.npz", + "liver": "liver_lesion.npz", + "kidney": "kidney_lesion.npz", + } + for organ, filename in lesion_files.items(): + path = os.path.join(seg_dir, filename) + if os.path.exists(path): + data = np.load(path)["data"] + voxels = int(np.sum(data > 0)) + if voxels > 0: + lesion_volume = round(voxels * voxel_volume, 2) + lesions[organ] = {"voxels": voxels, "volume": lesion_volume} + + organ_data_str = "" + for organ, vals in organ_volumes.items(): + organ_data_str += f"{organ.replace('_', ' ')}: volume={vals['volume']}cc, mean HU={vals['mean_hu']}\n" + + # If we have a RadGPT report, it is the authoritative source for organ status. + # Reset everything to normal first, then flag only what RadGPT calls abnormal. + if radgpt_comments: + for organ in list(organ_volumes.keys()): + organ_volumes[organ]['status'] = 'normal' + abnormal_keywords = ['enlarged', 'mass', 'lesion', 'tumor', 'abnormal', + 'dilated', 'obstruction', 'isoattenuating', 'hypodense', + 'hyperdense', 'cyst', 'nodule', 'atrophy', 'bilateral'] + # Build stripped root for flexible matching + organ_roots = {} + for organ in organ_volumes.keys(): + root = organ.replace('_left','').replace('_right','').replace('_body','') \ + .replace('_head','').replace('_tail','').replace('_gland','') \ + .replace('_duct','').replace('_lesion','').replace('_','') + organ_roots[organ] = root.lower() + for line in radgpt_comments.split('\n'): + line_stripped = line.lower().replace(' ','').replace('_','') + if any(kw in line.lower() for kw in abnormal_keywords): + for organ, root in organ_roots.items(): + # Skip subtypes unless explicitly mentioned + # e.g. "pancreas enlarged" shouldn't flag pancreas_body/head/tail + if '_body' in organ or '_head' in organ or '_tail' in organ or '_duct' in organ: + # Only flag subtype if the subtype word is in the line + subtype = organ.split('_')[-1] + if root in line_stripped and subtype in line.lower(): + organ_volumes[organ]['status'] = 'check' + else: + if root in line_stripped: + organ_volumes[organ]['status'] = 'check' + comments = radgpt_comments or "Clinical comments unavailable." + impression_items = radgpt_impression or ["No impression available for this case."] + + return jsonify({ + "case_id": id, + "patient": {"age": age, "sex": sex}, + "imaging": { + "study_type": study_detail, + "contrast": contrast, + "spacing": [round(float(s), 3) for s in spacing], + "shape": list(shape), + }, + "organ_volumes": organ_volumes, + "lesions": lesions, + "comments": comments, + "impression": impression_items, + }) + + except Exception as e: + import traceback + traceback.print_exc() + return jsonify({"error": "An internal error occurred."}), 500 + + @api_blueprint.route('/get-specific-segmentations/', methods=['POST']) async def get_specific_segmentations(combined_labels_id): combined_labels_id = combined_labels_id.replace("PanTS_", "") @@ -421,42 +749,52 @@ async def get_specific_segmentations(combined_labels_id): return response except Exception as e: return jsonify({"error": f"Error loading organ metrics: {str(e)}"}), 500 -@api_blueprint.route('/get-segmentations/', methods=['GET']) +@api_blueprint.route('/get-segmentations/.nii.gz', methods=['GET']) async def get_segmentations(combined_labels_id): - nifti_path = f"{Constants.PANTS_PATH}/mask_only/{get_panTS_id(combined_labels_id)}/{Constants.COMBINED_LABELS_NIFTI_FILENAME}" + if not _is_safe_id(combined_labels_id): + return jsonify({"error": "Invalid id"}), 400 + nifti_path = f"{Constants.PANTS_PATH}/mask_only/{get_panTS_id(secure_filename(combined_labels_id))}/{Constants.COMBINED_LABELS_NIFTI_FILENAME}" labels = list(Constants.PREDEFINED_LABELS.values()) - print("xdjs") # ?res=low → serve the precomputed low-res mask (paired with the low-res CT so the - # overlay stays aligned). Falls back to full res below if it hasn't been generated. + # overlay stays aligned). It lives under LOWRES_ROOT (a writable disk), NOT the + # read-only mask_only mount. Falls back to full res below if it hasn't been generated. if (request.args.get('res') or '').strip().lower() == 'low': - low_path = nifti_path.replace('.nii.gz', '_lowres.nii.gz') + low_name = Constants.COMBINED_LABELS_NIFTI_FILENAME.replace('.nii.gz', '_lowres.nii.gz') + low_path = f"{LOWRES_ROOT}/mask_only/{get_panTS_id(secure_filename(combined_labels_id))}/{low_name}" if os.path.exists(low_path): response = make_response(send_file(low_path, mimetype='application/gzip')) response.headers["Cross-Origin-Resource-Policy"] = "cross-origin" - response.headers['Content-Encoding'] = 'gzip' response.headers['Cache-Control'] = 'public, max-age=604800, immutable' return response img = nib.load(nifti_path) - print("⚠️ Detected float label map, converting to uint8 for Cornerstone compatibility...") try: + serve_path = nifti_path if img.get_data_dtype() != np.uint8: - raw = np.asanyarray(img.dataobj) - - rounded = np.rint(raw) - data = rounded.astype(np.uint8) - data = data.astype(np.uint8) - - new_img = nib.Nifti1Image(data, img.affine, header=img.header) - new_img.set_data_dtype(np.uint8) - - converted_path = nifti_path - nib.save(new_img, converted_path) - - response = make_response(send_file(nifti_path, mimetype='application/gzip')) + # The source is a float label map; Cornerstone needs uint8. Cache the + # converted copy in a WRITABLE temp dir and serve THAT — never write + # into the dataset's mask_only/ (it is read-only on the server, and an + # HTTP GET must not mutate ground-truth data). Writing the sibling into + # mask_only/ 500'd the segmentation endpoint in production. + cache_dir = "/tmp/pants_uint8" + os.makedirs(cache_dir, exist_ok=True) + converted_path = os.path.join( + cache_dir, + f"{get_panTS_id(secure_filename(combined_labels_id))}_combined_labels_uint8.nii.gz", + ) + if not os.path.exists(converted_path): + print("⚠️ Detected float label map, converting to uint8 for Cornerstone compatibility...") + raw = np.asanyarray(img.dataobj) + data = np.rint(raw).astype(np.uint8) + + new_img = nib.Nifti1Image(data, img.affine, header=img.header) + new_img.set_data_dtype(np.uint8) + nib.save(new_img, converted_path) + serve_path = converted_path + + response = make_response(send_file(serve_path, mimetype='application/gzip')) response.headers["Cross-Origin-Resource-Policy"] = "cross-origin" - response.headers['Content-Encoding'] = 'gzip' response.headers['Cache-Control'] = 'public, max-age=604800, immutable' # response.headers['Cross-Origin-Opener-Policy'] = 'same-origin' # response.headers['Cross-Origin-Embedder-Policy'] = 'require-corp' @@ -471,7 +809,9 @@ async def get_segmentations(combined_labels_id): @api_blueprint.route('/download/', methods=['GET']) def download_segmentation_zip(id): try: - outputs_ct_folder = Path(f"{Constants.PANTS_PATH}/mask_only/{get_panTS_id(id)}/segmentations") + if not _is_safe_id(id): + return jsonify({"error": "Invalid id"}), 400 + outputs_ct_folder = Path(f"{Constants.PANTS_PATH}/mask_only/{get_panTS_id(secure_filename(str(id)))}/segmentations") if not os.path.exists(outputs_ct_folder): return jsonify({"error": "Outputs/ct folder not found"}), 404 @@ -498,19 +838,80 @@ def download_segmentation_zip(id): print(f"❌ [Download Error] {e}") return jsonify({"error": "Internal server error"}), 500 -import threading import time inference_jobs = {} # {session_id: {status, model, error, session_path, zip_path}} +# Guards the read-modify-write in _set_inference_job: background segmentation +# threads and request handlers touch inference_jobs concurrently, and the +# get/update/set below is not atomic without a lock (updates would be lost). +_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): - current = inference_jobs.get(session_id, {}) - current.update(kwargs) - inference_jobs[session_id] = current + 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): +def _start_auto_segmentation(session_id, model_name, ct_file=None, server_input_path=None, precision="fp16"): + if not _is_safe_id(session_id): + return jsonify({"error": "Invalid session ID"}), 400 session_path = os.path.join(SESSIONS_DIR, session_id) os.makedirs(session_path, exist_ok=True) @@ -531,9 +932,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, @@ -541,9 +945,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, precision=precision, + ) + + 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}" @@ -569,6 +991,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)) @@ -622,6 +1049,11 @@ def _pick_text(*keys): session_id = _pick_text("session_id", "SESSION_ID", "sessionId") or str(uuid.uuid4()) model_name = _pick_text("model_name", "model", "MODEL_NAME") or "ePAI" + # Weight precision for the GPU inference stage. "fp16" routes to the quantized + # half-precision weights; "fp32" keeps the original full-precision checkpoints. + precision = (_pick_text("precision", "PRECISION") or "fp16").lower() + if precision not in ("fp16", "fp32"): + precision = "fp16" uploaded_filename = _pick_text("uploaded_filename", "output_filename", "filename") input_server_path = _pick_text("INPUT_SERVER_PATH", "input_server_path", "server_path", "path") source_reconstruction_session_id = _pick_text("source_reconstruction_session_id") @@ -633,7 +1065,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") @@ -682,17 +1114,148 @@ def _pick_text(*keys): model_name=model_name, ct_file=ct_file, server_input_path=input_server_path, + precision=precision, ) +# Catalog used by the automatic model picker. Each entry pairs a model with the +# tasks it is best suited for and its recommended inference precision. All of the +# segmentation networks run in FP16 by default (validated to produce masks +# identical to FP32); only the CPU-only ShapeKit stage is excluded. +MODEL_CATALOG = [ + {"model": "ePAI", "precision": "fp16", + "best_for": "pancreas, pancreatic duct, and pancreatic lesions (PDAC, cysts, PNETs)"}, + {"model": "SuPreM", "precision": "fp16", + "best_for": "broad multi-organ abdominal segmentation (25 structures)"}, + {"model": "MedFormer", "precision": "fp16", + "best_for": "26 abdominal structures plus pancreatic lesions; long-range context"}, + {"model": "R-Super", "precision": "fp16", + "best_for": "report-supervised segmentation of underrepresented structures"}, + {"model": "Atlas-Net", "precision": "fp16", + "best_for": "robust segmentation across varied CT protocols using anatomical priors"}, +] + +_VALID_MODELS = {m["model"] for m in MODEL_CATALOG} + + +def _heuristic_model_pick(hint: str): + """Rule-based fallback when the local LLM is unavailable.""" + text = (hint or "").lower() + if any(k in text for k in ("pancrea", "pdac", "cyst", "pnet", "duct", "lesion")): + choice = "ePAI" + elif any(k in text for k in ("organ", "abdomen", "abdominal", "liver", "kidney", + "spleen", "stomach", "multi", "whole", "full")): + choice = "SuPreM" + elif any(k in text for k in ("protocol", "robust", "varied", "generali")): + choice = "Atlas-Net" + else: + choice = "ePAI" + return { + "model": choice, + "precision": "fp16", + "reason": f"Rule-based match on your description; {choice} is best for " + + next(m["best_for"] for m in MODEL_CATALOG if m["model"] == choice) + + ". Running in FP16 (quantized) for speed.", + "source": "heuristic", + } + + +@api_blueprint.route('/suggest-model', methods=['POST']) +def suggest_model(): + """ + Automatically choose the best segmentation model (and precision) for a scan. + + Uses the local Ollama LLM to reason over the model catalog given a short, + optional free-text hint (e.g. "looking for pancreatic tumors") and/or the + uploaded filename. Falls back to a deterministic heuristic when the LLM is + unavailable, so the picker always returns a usable answer. + """ + payload = request.get_json(silent=True) or {} + hint = str(payload.get("hint") or request.form.get("hint") or "").strip() + filename = str(payload.get("filename") or request.form.get("filename") or "").strip() + combined_hint = " ".join(x for x in (hint, filename) if x) + + catalog_lines = "\n".join( + f"- {m['model']}: best for {m['best_for']}" for m in MODEL_CATALOG + ) + system_prompt = ( + "You are a routing assistant for a medical CT segmentation platform. " + "Choose exactly one model from the catalog that best fits the user's need. " + "Every model runs in FP16 (quantized) by default, which is faster and uses " + "half the memory with masks identical to FP32; only choose fp32 if the user " + "explicitly asks for maximum precision. Respond with a single JSON object: " + '{"model": , "precision": "fp16" or "fp32", ' + '"reason": }.' + ) + user_prompt = ( + f"Catalog:\n{catalog_lines}\n\n" + f"User need: {combined_hint or 'No specific description given; pick a sensible default.'}" + ) + + try: + result = chat_json( + model=DEFAULT_OLLAMA_MODEL, + system_prompt=system_prompt, + user_prompt=user_prompt, + temperature=0.1, + ) + model_choice = str(result.get("model", "")).strip() + if model_choice not in _VALID_MODELS: + return jsonify(_heuristic_model_pick(combined_hint)), 200 + precision = str(result.get("precision", "fp16")).strip().lower() + if precision not in ("fp16", "fp32"): + precision = "fp16" + reason = str(result.get("reason", "")).strip() or ( + f"{model_choice} is well suited to this task." + ) + return jsonify({ + "model": model_choice, + "precision": precision, + "reason": reason, + "source": "llm", + }), 200 + except (OllamaUnavailable, ValueError, Exception): + # Any LLM problem degrades gracefully to the deterministic heuristic. + return jsonify(_heuristic_model_pick(combined_hint)), 200 + + @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") @@ -909,10 +1472,17 @@ def download_pull_job_result(job_id): @api_blueprint.route('/get_result/', methods=['GET']) def get_result(session_id): + if not _is_safe_id(session_id): + return jsonify({"error": "Invalid session ID"}), 400 session_path = os.path.join(SESSIONS_DIR, session_id) zip_path = os.path.join(session_path, "auto_masks.zip") - wait_for_file(zip_path, timeout=30) + # Poll briefly for the archive. If it isn't ready, return 202 so the client + # can keep polling instead of receiving an uncaught TimeoutError → 500. + try: + wait_for_file(zip_path, timeout=30) + except TimeoutError: + return jsonify({"error": "Result not ready", "session_id": session_id}), 202 response = send_file( zip_path, @@ -925,7 +1495,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 @@ -937,7 +1507,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 @@ -953,7 +1523,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 @@ -996,6 +1566,13 @@ def upload_inference_chunk(): if not all([session_id, chunk_index, total_chunks, chunk_file]): return jsonify({"error": "Missing parameters"}), 400 + # session_id and chunk_index are both joined into a filesystem path below; + # reject anything non-numeric / traversal-y before it touches os.path. + if not _is_safe_id(session_id): + return jsonify({"error": "Invalid session ID"}), 400 + if not str(chunk_index).isdigit(): + return jsonify({"error": "Invalid chunk index"}), 400 + session_folder = os.path.join(CHUNK_DIR, session_id) os.makedirs(session_folder, exist_ok=True) @@ -1019,6 +1596,8 @@ def finalize_upload(): """ try: session_id = request.form.get("session_id") + if not _is_safe_id(session_id): + return jsonify({"error": "Invalid session ID"}), 400 total_chunks = int(request.form.get("total_chunks")) output_filename = request.form.get("output_filename", "inference_input.gz") requested_bdmap_id = request.form.get("bdmap_id") or request.form.get("case_id") @@ -1080,8 +1659,94 @@ def finalize_upload(): print(f"❌ Finalize upload error: {e}") return jsonify({"error": str(e)}), 500 +@api_blueprint.route("/upload-dicom-slice", methods=["POST"]) +def upload_dicom_slice(): + """Save a single DICOM slice to a session-specific temp directory.""" + try: + session_id = request.form.get("session_id") + if not session_id: + return jsonify({"error": "session_id required"}), 400 + slice_file = request.files.get("file") + if not slice_file: + return jsonify({"error": "file required"}), 400 + + dicom_dir = os.path.join("/tmp/uploads", session_id, "dicom") + os.makedirs(dicom_dir, exist_ok=True) + save_path = os.path.join(dicom_dir, slice_file.filename or f"{uuid.uuid4()}.dcm") + slice_file.save(save_path) + return jsonify({"status": "ok", "filename": os.path.basename(save_path)}) + except Exception as e: + print(f"❌ DICOM slice upload error: {e}") + return jsonify({"error": str(e)}), 500 + + +@api_blueprint.route("/finalize-dicom", methods=["POST"]) +def finalize_dicom(): + """Convert an uploaded DICOM series to NIfTI using SimpleITK.""" + try: + import SimpleITK as sitk + + session_id = request.form.get("session_id") + if not session_id: + return jsonify({"error": "session_id required"}), 400 + + dicom_dir = os.path.join("/tmp/uploads", session_id, "dicom") + if not os.path.isdir(dicom_dir): + return jsonify({"error": "No DICOM slices found for this session"}), 400 + + reader = sitk.ImageSeriesReader() + series_ids = reader.GetGDCMSeriesIDs(dicom_dir) + if not series_ids: + return jsonify({"error": "No valid DICOM series found in uploaded files"}), 400 + + dicom_names = reader.GetGDCMSeriesFileNames(dicom_dir, series_ids[0]) + reader.SetFileNames(dicom_names) + image = reader.Execute() + image = sitk.DICOMOrient(image, "LPS") + + # Save to sessions/inference///ct.nii.gz + digits = "".join(ch for ch in (session_id or "") if ch.isdigit()) + if len(digits) < 8: + fallback = f"{(uuid.uuid5(uuid.NAMESPACE_DNS, session_id).int % (10 ** 8)):08d}" + digits = (digits + fallback)[:8] + else: + digits = digits[:8] + bdmap_id = f"BDMAP_{digits}" + + base_path = os.path.join(Constants.SESSIONS_DIR_NAME, "inference", session_id) + target_dir = os.path.join(base_path, bdmap_id) + os.makedirs(target_dir, exist_ok=True) + final_path = os.path.join(target_dir, "ct.nii.gz") + + sitk.WriteImage(image, final_path) + + # Clean up temp DICOM slices + import shutil + shutil.rmtree(os.path.join("/tmp/uploads", session_id), ignore_errors=True) + + uploaded_filename = os.path.relpath(final_path, base_path) + return jsonify({ + "status": "converted", + "path": final_path, + "bdmap_id": bdmap_id, + "uploaded_filename": uploaded_filename, + }) + except Exception as e: + print(f"❌ DICOM finalize error: {e}") + return jsonify({"error": str(e)}), 500 + + ## OTHER ENDPOINTS ## +@api_blueprint.route('/cancel-inference', methods=['POST']) +def cancel_inference(): + cancel_all_inference() + for session_id, job in inference_jobs.items(): + if job.get('status') == 'running': + _set_inference_job(session_id, status='failed', error='Cancelled by user') + return jsonify({"message": "Inference cancelled"}), 200 + + @api_blueprint.route('/ping', methods=['GET']) def ping(): return jsonify({"message": "pong"}), 200 @@ -1090,8 +1755,6 @@ def ping(): def api_search(): # return jsonify({"message": "pong"}), 200 df = apply_filters(DF).copy() - sort_by = (_arg("sort_by", "top") or "top").strip().lower() - sort_by = (_arg("sort_by", "top") or "top").strip().lower() df = ensure_sort_cols(df) # ---- 排序參數 ---- @@ -1377,3 +2040,1632 @@ def api_random_topk_rotate_norand(): except Exception as e: return jsonify({"error": str(e)}), 400 + +AI_ALLOWED_ACTION_TYPES = { + "isolate_organs", + "show_organs", + "hide_organs", + "focus_organ", + "get_organ_metric", + "set_opacity", + "set_window", + "set_window_preset", + "set_zoom", + "zoom_to_fit", + "set_view", + "activate_measurement_tool", + "clear_measurements", + "list_structures", + "get_structure_count", + "get_largest_structure", + "get_smallest_structure", +} +AI_ALLOWED_VIEWS = {"mpr", "axial", "sagittal", "coronal", "3d"} +AI_ALLOWED_PRESETS = {"soft_tissue", "bone", "lung", "liver"} +AI_ALLOWED_TOOLS = {"distance", "probe", "roi"} +AI_ALLOWED_METRICS = {"volume_cm3", "mean_hu", "all"} + + +def _ai_norm(value): + return " ".join(str(value or "").lower().replace("_", " ").replace(".nii.gz", "").replace(".nii", "").split()) + + +def _ai_display(value): + text = str(value or "").replace(".nii.gz", "").replace(".nii", "").replace("_", " ").strip() + return text.title() if text else "Structure" + + +def _ai_metric_valid(value): + try: + number = float(value) + except (TypeError, ValueError): + return False + return np.isfinite(number) and number > 0 and number != 999999 + + +def _ai_public_metric(entry): + return { + "organ_name": str(entry.get("organ_name") or ""), + "display_name": _ai_display(entry.get("organ_name")), + "volume_cm3": entry.get("volume_cm3"), + "mean_hu": entry.get("mean_hu"), + } + + +def _ai_load_metrics(case_id, supplied_metrics): + """ + Prefer server-computed segmentation metrics. + + If the server does not have local NIfTI data, use metrics supplied by + the frontend. Never allow a missing local file to prevent Ollama from + answering general educational questions. + """ + + identifier = str(case_id or "").strip() + + if identifier and _is_safe_id(identifier): + try: + if identifier.isdigit(): + result = get_mask_data_internal(identifier) + source = "server_mask_data" + else: + result = calculate_session_metrics( + identifier, + Constants.SESSIONS_DIR_NAME, + ) + source = "session_mask_data" + + if isinstance(result, dict) and not result.get("error"): + raw_metrics = result.get("organ_metrics") or [] + + if isinstance(raw_metrics, list): + cleaned = [ + _ai_public_metric(item) + for item in raw_metrics + if isinstance(item, dict) + ] + + if cleaned: + return cleaned, source + + except Exception as error: + # Missing local case data should not disable general AI chat. + print( + "[AI metrics unavailable]", + type(error).__name__, + str(error), + ) + + if isinstance(supplied_metrics, list): + cleaned = [ + _ai_public_metric(item) + for item in supplied_metrics + if isinstance(item, dict) + ] + + if cleaned: + return cleaned, "frontend_supplied_metrics" + + return [], "unavailable" + + +def _ai_metric_lookup(metrics, available_organs): + lookup = {} + for entry in metrics: + organ_name = entry.get("organ_name") + display_name = entry.get("display_name") or _ai_display(organ_name) + for key in {organ_name, display_name, _ai_norm(organ_name), _ai_norm(display_name)}: + if key: + lookup[_ai_norm(key)] = entry + for organ in available_organs: + key = _ai_norm(organ) + if key not in lookup: + # Keep the viewer catalog entry available for honest “metric unavailable” responses. + lookup[key] = {"organ_name": organ, "display_name": _ai_display(organ), "volume_cm3": None, "mean_hu": None} + return lookup + + +def _ai_resolve_organ(value, available_organs): + norm = _ai_norm(value) + if not norm: + return None + + normalized_available = [(organ, _ai_norm(organ)) for organ in available_organs] + for organ, organ_norm in normalized_available: + if norm == organ_norm: + return organ + for organ, organ_norm in normalized_available: + if norm in organ_norm or organ_norm in norm: + return organ + + alias_groups = [ + {"left kidney", "kidney left", "left renal", "left renal kidney"}, + {"right kidney", "kidney right", "right renal", "right renal kidney"}, + {"gall bladder", "gallbladder"}, + {"inferior vena cava", "ivc", "vena cava", "postcava"}, + {"superior mesenteric artery", "sma"}, + {"common bile duct", "bile duct", "cbd"}, + {"left adrenal", "left adrenal gland", "adrenal gland left"}, + {"right adrenal", "right adrenal gland", "adrenal gland right"}, + ] + for aliases in alias_groups: + normalized_aliases = {_ai_norm(alias) for alias in aliases} + if norm not in normalized_aliases: + continue + for organ, organ_norm in normalized_available: + if organ_norm in normalized_aliases: + return organ + return None + + +def _ai_sanitize_actions(actions, available_organs): + if not isinstance(actions, list): + return [] + sanitized = [] + seen = set() + for action in actions[:8]: + if not isinstance(action, dict): + continue + action_type = action.get("type") + if action_type not in AI_ALLOWED_ACTION_TYPES: + continue + clean = {"type": action_type} + if action_type in {"isolate_organs", "show_organs", "hide_organs"}: + organs = action.get("organs") + if not isinstance(organs, list): + continue + resolved = [] + for organ in organs: + found = _ai_resolve_organ(organ, available_organs) + if found and found not in resolved: + resolved.append(found) + if not resolved: + continue + clean["organs"] = resolved + elif action_type == "focus_organ": + found = _ai_resolve_organ(action.get("organ"), available_organs) + if not found: + continue + clean["organ"] = found + elif action_type == "get_organ_metric": + found = _ai_resolve_organ(action.get("organ"), available_organs) + if not found: + continue + metric = action.get("metric") or "volume_cm3" + if metric not in AI_ALLOWED_METRICS: + metric = "volume_cm3" + clean["organ"] = found + clean["metric"] = metric + elif action_type == "set_opacity": + try: + clean["value"] = max(0, min(100, float(action.get("value")))) + except (TypeError, ValueError): + continue + elif action_type == "set_window": + try: + clean["width"] = max(1, float(action.get("width"))) + clean["center"] = float(action.get("center")) + except (TypeError, ValueError): + continue + elif action_type == "set_window_preset": + preset = action.get("preset") + if preset not in AI_ALLOWED_PRESETS: + continue + clean["preset"] = preset + elif action_type == "set_zoom": + try: + clean["value"] = max(0.1, min(20, float(action.get("value")))) + except (TypeError, ValueError): + continue + elif action_type == "set_view": + view = action.get("view") + if view not in AI_ALLOWED_VIEWS: + continue + clean["view"] = view + elif action_type == "activate_measurement_tool": + tool = action.get("tool") + if tool not in AI_ALLOWED_TOOLS: + continue + clean["tool"] = tool + + key = json.dumps(clean, sort_keys=True) + if key not in seen: + sanitized.append(clean) + seen.add(key) + return sanitized + + +def _ai_action_family(action_type): + if action_type in {"isolate_organs", "show_organs", "hide_organs", "focus_organ"}: + return "organ_visibility" + if action_type in {"set_window", "set_window_preset"}: + return "window" + if action_type in {"set_zoom", "zoom_to_fit"}: + return "zoom" + if action_type in {"activate_measurement_tool", "clear_measurements"}: + return "measurement_tool" + if action_type in {"list_structures", "get_structure_count", "get_largest_structure", "get_smallest_structure"}: + return "structure_query" + return action_type + + +def _ai_merge_actions(deterministic_actions, model_actions): + """Prefer deterministic actions for recognized commands and let Ollama fill gaps.""" + merged = [] + seen = set() + used_families = set() + for action in [*(deterministic_actions or []), *(model_actions or [])]: + if not isinstance(action, dict): + continue + key = json.dumps(action, sort_keys=True) + if key in seen: + continue + family = _ai_action_family(action.get("type")) + # A command should not apply conflicting visibility/window/zoom/tool actions. + if family in used_families and family != "get_organ_metric": + continue + merged.append(action) + seen.add(key) + used_families.add(family) + return merged[:8] + + +def _ai_metadata(case_id, supplied): + metadata = {} + if isinstance(supplied, dict): + for key in ["sex", "age", "bmi", "height_cm", "weight_kg"]: + if supplied.get(key) not in [None, ""]: + metadata[key] = supplied.get(key) + if case_id and str(case_id).isdigit(): + entry = _METADATA_CACHE.get(get_panTS_id(str(case_id)), {}) + if entry.get("age") not in [None, ""]: + try: + metadata["age"] = float(entry.get("age")) + except (TypeError, ValueError): + metadata["age"] = entry.get("age") + if entry.get("sex") not in [None, ""]: + metadata["sex"] = entry.get("sex") + if "bmi" not in metadata: + try: + height_m = float(metadata.get("height_cm")) / 100 + weight_kg = float(metadata.get("weight_kg")) + if height_m > 0 and weight_kg > 0: + metadata["bmi"] = round(weight_kg / (height_m * height_m), 1) + except (TypeError, ValueError): + pass + return metadata + + +def _ai_has_case_reference(norm: str) -> bool: + case_phrases = ( + "this scan", + "this ct", + "this case", + "this patient", + "this segmentation", + "the segmentation", + "current scan", + "current case", + "currently loaded", + "shown here", + "in the viewer", + "in this image", + "in these images", + "my scan", + "my ct", + "my liver", + "my pancreas", + "my kidney", + "my spleen", + "do i have", + "am i", + "for this patient", + "patient s", + "patient's", + ) + + return any(phrase in norm for phrase in case_phrases) + + +def _ai_question_mode(message, fallback_actions): + """ + Separate general knowledge from case-specific questions. + + Modes: + - general_education + - case_metadata + - case_measurement + - case_health_context + - viewer_command + """ + + norm = _ai_norm(message) + action_types = { + action.get("type") + for action in (fallback_actions or []) + if isinstance(action, dict) + } + + viewer_action_types = { + "isolate_organs", + "show_organs", + "hide_organs", + "focus_organ", + "set_opacity", + "set_window", + "set_window_preset", + "set_zoom", + "zoom_to_fit", + "set_view", + "activate_measurement_tool", + "clear_measurements", + } + + if action_types.intersection(viewer_action_types): + return "viewer_command" + + has_case_reference = _ai_has_case_reference(norm) + + if "bmi" in norm or "body mass index" in norm: + case_bmi_phrases = ( + "patient bmi", + "patient s bmi", + "patient's bmi", + "this bmi", + "their bmi", + "his bmi", + "her bmi", + "my bmi", + "bmi for this", + "bmi of this", + ) + + if has_case_reference or any( + phrase in norm for phrase in case_bmi_phrases + ): + return "case_metadata" + + return "general_education" + + if "age" in norm or "how old" in norm: + if has_case_reference or "patient age" in norm: + return "case_metadata" + + health_terms = ( + "healthy", + "unhealthy", + "normal", + "abnormal", + "concerning", + "disease", + "diseased", + "cancer", + "cancerous", + "tumor", + "malignant", + "benign", + "enlarged", + "too large", + "too small", + "swollen", + "damaged", + "cirrhosis", + "fatty liver", + "lesion", + "mass", + ) + + health_question = any( + term in norm for term in health_terms + ) + + implied_current_case = ( + norm.startswith("is the liver ") + or norm.startswith("is the pancreas ") + or norm.startswith("is the spleen ") + or norm.startswith("is the kidney ") + or norm.startswith("does the liver look ") + ) + + if health_question and ( + has_case_reference or implied_current_case + ): + return "case_health_context" + + measurement_terms = ( + "volume", + "how big", + "what size", + "size of", + "mean hu", + "hounsfield", + "largest structure", + "smallest structure", + "how many structures", + "structure count", + "measured", + ) + + asks_measurement = any( + term in norm for term in measurement_terms + ) + + general_reference_terms = ( + "normal liver volume", + "normal organ volume", + "typical liver volume", + "average liver volume", + "usual liver size", + "what is considered normal", + ) + + if asks_measurement: + if ( + any(term in norm for term in general_reference_terms) + and not has_case_reference + ): + return "general_education" + + return "case_measurement" + + return "general_education" + + +def _ai_case_metadata_reply(norm, metadata): + """ + Answer only explicitly case-specific age/BMI questions. + + General questions such as 'What is BMI?' must continue to Ollama. + """ + + has_case_reference = _ai_has_case_reference(norm) + + asks_case_age = ( + ("age" in norm or "how old" in norm) + and ( + has_case_reference + or "patient age" in norm + or "age of the patient" in norm + ) + ) + + if asks_case_age: + if metadata.get("age") not in [None, ""]: + age = metadata.get("age") + + try: + age = round(float(age)) + except (TypeError, ValueError): + pass + + return ( + "The available metadata lists this patient's age " + f"as **{age} years**." + ) + + return ( + "The current case metadata does not include a valid age." + ) + + asks_case_bmi = ( + ("bmi" in norm or "body mass index" in norm) + and ( + has_case_reference + or "patient bmi" in norm + or "patient's bmi" in norm + or "bmi of the patient" in norm + or "bmi for the patient" in norm + ) + ) + + if asks_case_bmi: + if metadata.get("bmi") not in [None, ""]: + try: + bmi = float(metadata.get("bmi")) + return ( + "The available metadata lists this patient's BMI " + f"as **{bmi:.1f}**." + ) + except (TypeError, ValueError): + pass + + return ( + "The current case does not include enough height and weight " + "information to calculate the patient's BMI." + ) + + return None + + +def _ai_relevant_metrics(message, metrics): + norm = _ai_norm(message) + relevant = [] + + for entry in metrics or []: + if not isinstance(entry, dict): + continue + + organ_name = entry.get("organ_name") + display_name = ( + entry.get("display_name") + or _ai_display(organ_name) + ) + + names = { + _ai_norm(organ_name), + _ai_norm(display_name), + } + + if any(name and name in norm for name in names): + relevant.append(entry) + + return relevant + + +def _ai_case_facts(message, metrics, metadata): + """ + Produce exact deterministic facts that can be placed beside the model's + health-context explanation. + """ + + facts = [] + relevant_metrics = _ai_relevant_metrics( + message, + metrics, + ) + + for entry in relevant_metrics: + organ = ( + entry.get("display_name") + or _ai_display(entry.get("organ_name")) + ) + + volume = entry.get("volume_cm3") + mean_hu = entry.get("mean_hu") + + organ_facts = [] + + if _ai_metric_valid(volume): + organ_facts.append( + f"segmented volume {float(volume):.2f} cm³" + ) + + if _ai_metric_valid(mean_hu) or mean_hu == 0: + organ_facts.append( + f"mean attenuation {float(mean_hu):.1f} HU" + ) + + if organ_facts: + facts.append( + f"**{organ}:** " + ", ".join(organ_facts) + ) + + if metadata.get("age") not in [None, ""]: + facts.append(f"**Age:** {metadata.get('age')}") + + if metadata.get("sex") not in [None, ""]: + facts.append(f"**Sex:** {metadata.get('sex')}") + + if metadata.get("bmi") not in [None, ""]: + facts.append(f"**BMI:** {metadata.get('bmi')}") + + return facts + + +def _ai_extract_unknown_volume_request(norm, metric_lookup): + if "volume" not in norm and "how big" not in norm and "size" not in norm: + return None + for key, entry in metric_lookup.items(): + if key and key in norm: + return entry + # A small deny-list for common out-of-scope organs people may ask about. + for organ in ["brain", "heart", "eye", "skull"]: + if organ in norm: + return {"organ_name": organ, "display_name": _ai_display(organ), "volume_cm3": None, "mean_hu": None, "missing": True} + return None + + +def _ai_action_prefix(actions): + bits = [] + if any(a.get("type") == "set_view" and a.get("view") == "3d" for a in actions): + bits.append("switched to 3D") + isolate = next((a for a in actions if a.get("type") == "isolate_organs"), None) + if isolate: + bits.append("isolated " + ", ".join(_ai_display(o) for o in isolate.get("organs", []))) + elif any(a.get("type") == "show_organs" for a in actions): + show = next(a for a in actions if a.get("type") == "show_organs") + bits.append("showed " + ", ".join(_ai_display(o) for o in show.get("organs", []))) + if any(a.get("type") == "focus_organ" for a in actions): + focus = next(a for a in actions if a.get("type") == "focus_organ") + bits.append("focused on " + _ai_display(focus.get("organ"))) + if bits: + if len(bits) == 1: + return "I " + bits[0] + ". " + return "I " + ", ".join(bits[:-1]) + ", and " + bits[-1] + ". " + return "" + + +def _ai_structure_names(metrics, available_organs): + raw_names = [m.get("organ_name") for m in metrics if m.get("organ_name")] + if not raw_names: + raw_names = available_organs + names = [] + seen = set() + for name in raw_names: + display = _ai_display(name) + key = _ai_norm(display) + if key and key not in seen: + names.append(display) + seen.add(key) + return names + + +def _ai_action_confirmation(actions): + confirmations = [] + for action in actions: + action_type = action.get("type") + if action_type == "isolate_organs": + confirmations.append("Isolated " + ", ".join(_ai_display(o) for o in action.get("organs", [])) + ".") + elif action_type == "show_organs": + confirmations.append("Showed " + ", ".join(_ai_display(o) for o in action.get("organs", [])) + ".") + elif action_type == "hide_organs": + confirmations.append("Hid " + ", ".join(_ai_display(o) for o in action.get("organs", [])) + ".") + elif action_type == "focus_organ": + confirmations.append("Focused on " + _ai_display(action.get("organ")) + ".") + elif action_type == "set_view": + confirmations.append(f"Switched to {str(action.get('view')).upper()} view.") + elif action_type == "set_opacity": + confirmations.append(f"Set overlay opacity to {float(action.get('value')):.0f}%.") + elif action_type == "set_window_preset": + confirmations.append(f"Applied the {str(action.get('preset')).replace('_', ' ')} window preset.") + elif action_type == "set_window": + confirmations.append(f"Set the CT window to width {float(action.get('width')):.0f} and center {float(action.get('center')):.0f}.") + elif action_type == "set_zoom": + confirmations.append(f"Set zoom to {float(action.get('value')):g}.") + elif action_type == "zoom_to_fit": + confirmations.append("Reset zoom to fit.") + elif action_type == "activate_measurement_tool": + confirmations.append(f"Activated the {str(action.get('tool')).upper()} tool.") + elif action_type == "clear_measurements": + confirmations.append("Cleared the current measurements.") + return " ".join(confirmations) + + +def _ai_grounded_reply( + message, + actions, + metrics, + available_organs, + metadata, + candidate_reply, + question_mode, +): + norm = _ai_norm(message) + metric_lookup = _ai_metric_lookup( + metrics, + available_organs, + ) + + # Only intercept age/BMI when the user explicitly asks about + # the currently loaded patient. + if question_mode == "case_metadata": + metadata_reply = _ai_case_metadata_reply( + norm, + metadata, + ) + + if metadata_reply: + return metadata_reply + + # For health questions, preserve Ollama's explanation while placing + # exact measured case facts above it. + if question_mode == "case_health_context": + facts = _ai_case_facts( + message, + metrics, + metadata, + ) + + if ( + isinstance(candidate_reply, str) + and candidate_reply.strip() + ): + explanation = candidate_reply.strip() + else: + explanation = ( + "The available segmentation measurements can provide " + "useful context, but they are not enough by themselves " + "to determine whether this organ is healthy. A complete " + "assessment would also consider the CT appearance, " + "attenuation, enhancement, contour, focal lesions, " + "clinical history, and relevant laboratory results." + ) + + if facts: + reply = ( + "Available measurements for this case:\n" + + "\n".join(f"- {fact}" for fact in facts) + + "\n\n" + + explanation + ) + else: + reply = explanation + + uncertainty_terms = ( + "not a diagnosis", + "cannot confirm", + "cannot determine", + "volume alone", + "measurements alone", + "radiologist", + "clinical evaluation", + ) + + if not any( + term in explanation.lower() + for term in uncertainty_terms + ): + reply += ( + "\n\nThis is an educational, non-diagnostic assessment. " + "Volume and segmentation measurements alone cannot " + "confirm that an organ is healthy or diagnose disease." + ) + + confirmation = _ai_action_confirmation(actions) + + if confirmation: + reply += "\n\n" + confirmation + + return reply + + # Deterministically ground exact case measurements. + for action in actions: + if action.get("type") != "get_organ_metric": + continue + + entry = metric_lookup.get( + _ai_norm(action.get("organ")) + ) + + organ_label = _ai_display( + action.get("organ") + ) + + if not entry or entry.get("missing"): + return ( + f"{organ_label} is not available in this segmentation, " + f"so no {organ_label.lower()} measurement is available " + "for this case." + ) + + metric = action.get("metric") or "volume_cm3" + volume = entry.get("volume_cm3") + mean_hu = entry.get("mean_hu") + prefix = _ai_action_prefix(actions) + + if metric == "mean_hu": + if _ai_metric_valid(mean_hu) or mean_hu == 0: + return ( + f"{prefix}The segmented **{organ_label}** mean " + f"attenuation is **{float(mean_hu):.1f} HU**, " + "measured from the segmentation mask." + ) + + return ( + f"{prefix}Mean HU for **{organ_label}** is not " + "available for this case." + ) + + if metric == "all": + parts = [] + + if _ai_metric_valid(volume): + parts.append( + f"volume **{float(volume):.2f} cm³**" + ) + + if _ai_metric_valid(mean_hu) or mean_hu == 0: + parts.append( + f"mean attenuation " + f"**{float(mean_hu):.1f} HU**" + ) + + if parts: + return ( + f"{prefix}For the segmented **{organ_label}**, " + + " and ".join(parts) + + "." + ) + + return ( + f"{prefix}Metrics for **{organ_label}** are not " + "available for this case." + ) + + if _ai_metric_valid(volume): + return ( + f"{prefix}The segmented **{organ_label}** volume is " + f"**{float(volume):.2f} cm³**, measured from the " + "segmentation mask." + ) + + return ( + f"{prefix}No valid segmented volume is available for " + f"**{organ_label}** in this case." + ) + + if any( + action.get("type") == "get_largest_structure" + for action in actions + ): + valid = [ + metric + for metric in metrics + if _ai_metric_valid(metric.get("volume_cm3")) + ] + + if valid: + largest = max( + valid, + key=lambda item: float( + item.get("volume_cm3") + ), + ) + + return ( + "The largest segmented structure is " + f"**{_ai_display(largest.get('organ_name'))}**, " + "with a measured volume of " + f"**{float(largest.get('volume_cm3')):.2f} cm³**." + ) + + return ( + "I could not determine the largest structure because valid " + "volume metrics are unavailable for this case." + ) + + if any( + action.get("type") == "get_smallest_structure" + for action in actions + ): + valid = [ + metric + for metric in metrics + if _ai_metric_valid(metric.get("volume_cm3")) + ] + + if valid: + smallest = min( + valid, + key=lambda item: float( + item.get("volume_cm3") + ), + ) + + return ( + "The smallest segmented structure is " + f"**{_ai_display(smallest.get('organ_name'))}**, " + "with a measured volume of " + f"**{float(smallest.get('volume_cm3')):.2f} cm³**." + ) + + return ( + "I could not determine the smallest structure because valid " + "volume metrics are unavailable for this case." + ) + + if any( + action.get("type") == "get_structure_count" + for action in actions + ): + names = _ai_structure_names( + metrics, + available_organs, + ) + + return ( + f"This case has **{len(names)} segmented structures** " + "listed for the viewer." + ) + + if any( + action.get("type") == "list_structures" + for action in actions + ): + names = _ai_structure_names( + metrics, + available_organs, + ) + + if names: + return ( + f"This case includes **{len(names)} segmented " + "structures**: " + + ", ".join(names) + + "." + ) + + return ( + "No segmented structures are listed for this case." + ) + + unknown_volume = _ai_extract_unknown_volume_request( + norm, + metric_lookup, + ) + + if unknown_volume and unknown_volume.get("missing"): + organ_name = _ai_display( + unknown_volume.get("organ_name") + ) + + return ( + f"The **{organ_name}** is not included in this abdominal " + f"segmentation, so no {organ_name.lower()} volume is " + "available for this case." + ) + + confirmation = _ai_action_confirmation(actions) + + if ( + isinstance(candidate_reply, str) + and candidate_reply.strip() + ): + reply = candidate_reply.strip() + + if confirmation: + reply += "\n\n" + confirmation + + return reply + + if confirmation: + return confirmation + + return ( + "I could not generate a complete response. Please verify that " + "Ollama is running and that the selected model is installed." + ) + + +def _ai_system_prompt(): + return """ +You are BodyMaps AI, a capable and conversational assistant embedded in +an abdominal CT visualization application. + +Your job is to answer the user's actual question. Do not assume that +every question asks about the current patient. + +You operate in several modes. + +GENERAL EDUCATIONAL QUESTIONS +Answer ordinary questions using your general knowledge. + +Examples: +- What is a CT scan? +- What is body mass index? +- What does the liver do? +- What are Hounsfield units? +- What is a segmentation mask? +- What conditions can affect the liver? +- What is a typical liver volume? + +These questions do not require patient metadata. Answer them naturally, +clearly, and conversationally. + +CASE-SPECIFIC MEASUREMENTS +When the user asks about this scan, this case, this patient, the current +segmentation, or a measured structure, use the supplied case data. + +Examples: +- What is the liver volume in this scan? +- How big is the segmented pancreas? +- What is this patient's BMI? +- Which segmented structure is largest? + +Never invent a case-specific value. Use exact supplied values. + +CASE-SPECIFIC HEALTH QUESTIONS +When the user asks questions such as: +- Is this liver healthy? +- Does this liver appear enlarged? +- Is this organ normal? +- Is this finding concerning? +- Could this be a tumor? + +Provide a useful evidence-limited assessment rather than a generic +refusal. + +You should: +1. State the exact available case measurements. +2. Explain what those measurements may suggest. +3. Compare them with broad educational expectations when appropriate, + while clearly labeling those comparisons as approximate. +4. Explain important limitations of the available evidence. +5. State what additional imaging features, metadata, laboratory values, + symptoms, or clinical history would normally be considered. + +Do not state that a patient definitely has or does not have a disease +when the supplied data does not establish that conclusion. + +Volume alone does not prove that an organ is healthy or unhealthy. +A segmented volume can provide useful context, but complete assessment +may also require contour, attenuation, contrast enhancement, focal +lesions, surrounding structures, prior scans, laboratory results, and +clinical history. + +You may explain possible diseases and general treatment approaches. +Do not prescribe medication, choose a personalized treatment, or tell +the user to ignore professional medical care. + +VIEWER COMMANDS +Translate viewer requests into the allowed structured actions. +Actions are applied immediately and should not require confirmation. + +GROUNDING RULES +- Never invent organ volume, mean HU, age, BMI, or patient metadata. +- General medical and imaging knowledge does not need to be present in + the case metadata. +- Patient-specific facts must come from supplied data. +- If case measurements are unavailable, say which data is missing and + still answer the educational portion of the question. +- "Segment the liver" means display or isolate an existing segmentation. + Do not claim a new segmentation model was run unless the backend + actually reports that it ran. + +OUTPUT FORMAT +Return exactly one JSON object: + +{ + "reply": "complete conversational answer", + "actions": [], + "intent": "short_intent_name" +} + +Allowed actions: +- {"type":"isolate_organs","organs":["exact available organ name"]} +- {"type":"show_organs","organs":["exact available organ name"]} +- {"type":"hide_organs","organs":["exact available organ name"]} +- {"type":"focus_organ","organ":"exact available organ name"} +- {"type":"get_organ_metric","organ":"exact available organ name","metric":"volume_cm3|mean_hu|all"} +- {"type":"list_structures"} +- {"type":"get_structure_count"} +- {"type":"get_largest_structure"} +- {"type":"get_smallest_structure"} +- {"type":"set_view","view":"mpr|axial|sagittal|coronal|3d"} +- {"type":"set_opacity","value":0-100} +- {"type":"set_window","width":number,"center":number} +- {"type":"set_window_preset","preset":"soft_tissue|bone|lung|liver"} +- {"type":"set_zoom","value":number} +- {"type":"zoom_to_fit"} +- {"type":"activate_measurement_tool","tool":"distance|probe|roi"} +- {"type":"clear_measurements"} + +For a question that does not need a viewer action, return an empty +actions list. + +Return JSON only. +""".strip() + + +@api_blueprint.route("/ai-models", methods=["GET"]) +def ai_models(): + try: + models = list_ollama_models() + model_names = [model["name"] for model in models] + default_model = DEFAULT_OLLAMA_MODEL if DEFAULT_OLLAMA_MODEL in model_names else (model_names[0] if model_names else DEFAULT_OLLAMA_MODEL) + return jsonify({ + "available": True, + "models": models, + "default_model": default_model, + }) + except OllamaUnavailable as error: + return jsonify({ + "available": False, + "models": [], + "default_model": DEFAULT_OLLAMA_MODEL, + "error": f"Ollama is not reachable at the configured local endpoint: {error}", + }), 200 + + +@api_blueprint.route("/ai-command", methods=["POST"]) +def ai_command(): + try: + body = request.get_json( + force=True, + silent=True, + ) or {} + + message = str( + body.get("message") or "" + ).strip() + + if not message: + return jsonify( + { + "reply": ( + "Please type a question or viewer command." + ), + "actions": [], + "source": "validation", + } + ), 400 + + available_organs = body.get( + "available_organs" + ) or [] + + if not isinstance(available_organs, list): + available_organs = [] + + available_organs = [ + str(item).strip() + for item in available_organs + if str(item).strip() + ] + + viewer_state = ( + body.get("viewer_state") + if isinstance( + body.get("viewer_state"), + dict, + ) + else {} + ) + + case_id = str( + body.get("session_id") + or body.get("case_id") + or "" + ).strip() + + requested_model = body.get("model") + + # Always use the configured default when the frontend does not + # explicitly send a model. + selected_model = ( + requested_model.strip() + if isinstance(requested_model, str) + and requested_model.strip() + else DEFAULT_OLLAMA_MODEL + ) + + metrics, metric_source = _ai_load_metrics( + case_id, + body.get("organ_metrics"), + ) + + metadata = _ai_metadata( + case_id, + body.get("demographics"), + ) + + # Include metric organ names even if the frontend organ catalog + # was empty or incomplete. + for metric in metrics: + organ_name = str( + metric.get("organ_name") or "" + ).strip() + + if ( + organ_name + and organ_name not in available_organs + ): + available_organs.append(organ_name) + + fallback = parse_intent( + message=message, + available_organs=available_organs, + viewer_state=viewer_state, + case_id=case_id or None, + ) + + fallback_actions = _ai_sanitize_actions( + fallback.get("actions", []), + available_organs, + ) + + question_mode = _ai_question_mode( + message, + fallback_actions, + ) + + # Do not use the fallback's medical refusal as the model's answer. + # The model receives the fallback only as an action suggestion. + rule_suggestion_reply = None + + if question_mode in { + "viewer_command", + "case_measurement", + }: + rule_suggestion_reply = fallback.get("reply") + + prompt_payload = { + "request_mode": question_mode, + "user_message": message, + "current_case": { + "case_id": case_id or None, + "available_organs": available_organs, + "computed_organ_metrics": metrics, + "metadata": metadata, + "metrics_source": metric_source, + }, + "viewer_state": viewer_state, + "rule_based_suggestion": { + "reply": rule_suggestion_reply, + "actions": fallback_actions, + "intent": fallback.get("intent"), + }, + "response_requirements": { + "answer_general_questions_using_model_knowledge": True, + "use_exact_case_values_when_case_specific": True, + "do_not_invent_patient_specific_values": True, + "provide_non_diagnostic_health_context": True, + "return_json_only": True, + }, + } + + model_result = None + model_error = None + source = "ollama" + + try: + model_result = chat_json( + model=selected_model, + system_prompt=_ai_system_prompt(), + user_prompt=json.dumps( + prompt_payload, + ensure_ascii=False, + ), + temperature=0.2, + ) + except ( + OllamaUnavailable, + Exception, + ) as error: + # Keep viewer actions usable if Ollama is temporarily offline. + model_error = str(error) + model_result = None + source = "rule_fallback" + + print( + "[Ollama unavailable]", + type(error).__name__, + model_error, + ) + + if isinstance(model_result, dict): + model_actions = _ai_sanitize_actions( + model_result.get("actions", []), + available_organs, + ) + + actions = _ai_merge_actions( + fallback_actions, + model_actions, + ) + + candidate_reply = ( + model_result.get("reply") + or rule_suggestion_reply + ) + + intent = ( + model_result.get("intent") + or fallback.get("intent") + or question_mode + ) + else: + actions = fallback_actions + candidate_reply = rule_suggestion_reply + intent = ( + fallback.get("intent") + or question_mode + ) + + reply = _ai_grounded_reply( + message=message, + actions=actions, + metrics=metrics, + available_organs=available_organs, + metadata=metadata, + candidate_reply=candidate_reply, + question_mode=question_mode, + ) + + response = { + "reply": reply, + "actions": actions, + "grounding": { + "case_id": case_id, + "request_mode": question_mode, + "metrics_source": metric_source, + "organ_count": ( + len(metrics) + if metrics + else len(available_organs) + ), + "metadata_fields": sorted( + metadata.keys() + ), + }, + "source": source, + "model": ( + selected_model + if source == "ollama" + else None + ), + "intent": intent, + } + + if model_error: + response["ollama_error"] = model_error + + return jsonify(response) + + except Exception as error: + print( + "[ai_command error]", + type(error).__name__, + str(error), + ) + + return jsonify( + { + "reply": ( + "An internal error occurred while processing " + "the AI request." + ), + "actions": [], + "source": "error", + "error_type": type(error).__name__, + } + ), 500 + +# --------------------------------------------------------------------------- +# Edited segmentation masks (viewer's Edit Masks panel). +# Strictly additive and isolated: writes ONLY into {PANTS_PATH}/edited_masks/, +# never touching image_only/ or mask_only/, so the original dataset is safe. +# Each save is timestamped rather than overwritten — a lightweight version +# history a maintainer can inspect or promote manually. +# --------------------------------------------------------------------------- + +_EDITED_MASKS_DIRNAME = "edited_masks" +_EDITED_MASK_MAX_BYTES = 512 * 1024 * 1024 # generous cap for a full-body labelmap + + +def _edited_masks_dir(case_id): + # Traversal safety: require digits, then convert to int before it reaches the + # filesystem. A number can't carry a "../" or "/" payload, so the only value + # that flows into os.path.join is fully controlled (get_panTS_id just zero-pads + # and prefixes "PanTS_"). The int() cast is also what lets static analysis + # (CodeQL py/path-injection) see the user-tainted string is neutralized. + if not str(case_id).isdigit(): + raise ValueError("case_id must be numeric") + return os.path.join(Constants.PANTS_PATH, _EDITED_MASKS_DIRNAME, get_panTS_id(int(case_id))) + + +@api_blueprint.route('/save-edited-mask/', methods=['POST']) +def save_edited_mask(case_id): + try: + uploaded = request.files.get("mask") + if uploaded is None: + return jsonify({"error": "No file field 'mask' in the request."}), 400 + data = uploaded.read(_EDITED_MASK_MAX_BYTES + 1) + if len(data) > _EDITED_MASK_MAX_BYTES: + return jsonify({"error": "Mask file too large."}), 413 + # The viewer always sends gzipped NIfTI; check the gzip magic bytes. + if len(data) < 2 or data[0] != 0x1F or data[1] != 0x8B: + return jsonify({"error": "Expected a gzipped NIfTI (.nii.gz) file."}), 400 + + out_dir = _edited_masks_dir(case_id) + os.makedirs(out_dir, exist_ok=True) + timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") + filename = f"combined_labels_edited_{timestamp}.nii.gz" + with open(os.path.join(out_dir, filename), "wb") as f: + f.write(data) + # Side car carrying custom class names/colors + # the .nii.gz only stores integer labels, so this saves user-created class names/colors + labels_field = request.files.get("labels") + if labels_field is not None: + labels_data = labels_field.read(1024 * 1024) + try: + parsed = json.loads(labels_data.decode("utf-8")) + except Exception: + parsed = None + if isinstance(parsed, dict): + labels_filename = filename.replace(".nii.gz", "_labels.json") + with open(os.path.join(out_dir, labels_filename), "w") as f: + json.dump(parsed, f) + + return jsonify({"saved": True, "filename": filename, "bytes": len(data)}) + except Exception as error: + print("[save_edited_mask error]", type(error).__name__, error) + return jsonify({"error": "Failed to save the edited mask."}), 500 + + +@api_blueprint.route('/list-edited-masks/', methods=['GET']) +def list_edited_masks(case_id): + try: + out_dir = _edited_masks_dir(case_id) + if not os.path.isdir(out_dir): + return jsonify({"items": []}) + items = [] + for name in sorted(os.listdir(out_dir), reverse=True): + path = os.path.join(out_dir, name) + if not os.path.isfile(path): + continue + items.append({ + "filename": name, + "bytes": os.path.getsize(path), + "modified": datetime.fromtimestamp(os.path.getmtime(path)).isoformat(), + }) + return jsonify({"items": items}) + except Exception as error: + print("[list_edited_masks error]", type(error).__name__, error) + return jsonify({"error": "Failed to list edited masks."}), 500 + + +# --------------------------------------------------------------------------- +# Advanced analysis (viewer's AI-segment tool + vessel CPR panel). +# Additive and read-only against the dataset — loads image_only/ and mask_only/ +# but writes nothing there. Heavy numeric work lives in services/advanced_analysis. +# --------------------------------------------------------------------------- + +import threading + +# Each of these requests loads a CT volume into worker RAM, so unbounded +# concurrency is an OOM waiting to happen. Cap in-flight analyses per process +# (gunicorn workers each get their own slots); extras get an immediate 503 +# instead of queueing until the worker starves. +_ANALYSIS_SLOTS = threading.BoundedSemaphore(2) +_ANALYSIS_BUSY_RESPONSE = ( + {"error": "The analysis service is busy — try again in a moment."}, + 503, +) + +def _safe_case_id(case_id): + # Traversal safety: require digits, then hand get_panTS_id an int so the + # user value can't carry a "../" or "/" payload into the CT/mask path. The + # int() cast is also the barrier CodeQL recognizes as sanitizing the taint. + if not str(case_id).isdigit(): + raise ValueError("case_id must be numeric") + return int(case_id) + + +def _case_ct_path(case_id, low=False): + case_dir = f"{Constants.PANTS_PATH}/image_only/{get_panTS_id(_safe_case_id(case_id))}" + path = f"{case_dir}/{Constants.MAIN_NIFTI_FILENAME}" + if low: + low_path = path.replace('.nii.gz', '_lowres.nii.gz') + if os.path.exists(low_path): + return low_path + return path + + +def _case_mask_path(case_id, low=False): + case_dir = f"{Constants.PANTS_PATH}/mask_only/{get_panTS_id(_safe_case_id(case_id))}" + path = f"{case_dir}/{Constants.COMBINED_LABELS_NIFTI_FILENAME}" + if low: + low_path = path.replace('.nii.gz', '_lowres.nii.gz') + if os.path.exists(low_path): + return low_path + return path + + +@api_blueprint.route('/interactive-segment/', methods=['POST']) +def interactive_segment(case_id): + """Click-to-segment: seed prompt -> proposed mask (.nii.gz in CT geometry). + + Body JSON: { point_lps:[x,y,z] | point_ijk:[i,j,k], tolerance?, box_lps?, + res?: "low"|"full" }. res should match the resolution the viewer + loaded so the returned mask's voxel grid aligns with the labelmap. + """ + if not _ANALYSIS_SLOTS.acquire(blocking=False): + return jsonify(_ANALYSIS_BUSY_RESPONSE[0]), _ANALYSIS_BUSY_RESPONSE[1] + try: + import numpy as np + from services.advanced_analysis import segment_from_prompt + body = request.get_json(force=True, silent=True) or {} + low = (body.get("res") or "low").lower() == "low" + ct_path = _case_ct_path(case_id, low=low) + if not os.path.exists(ct_path): + return jsonify({"error": "CT not found for this case on the server."}), 404 + + ct_obj = nib.load(ct_path) + # float32: half the RAM of nibabel's float64 default — these are public + # endpoints and a full-res CT at float64 is multiple GB per request. + ct = ct_obj.get_fdata(dtype=np.float32) + mask = segment_from_prompt(ct, ct_obj.affine, body) + if int(mask.sum()) == 0: + return jsonify({"error": "Nothing grew from that point — try a different spot or a higher tolerance."}), 422 + + out = nib.Nifti1Image(mask, ct_obj.affine, ct_obj.header) + out.header.set_data_dtype('uint8') + # nibabel serializes an uncompressed .nii to bytes; gzip it ourselves. + import gzip as _gzip + gz = _gzip.compress(out.to_bytes()) + resp = make_response(gz) + resp.headers['Content-Type'] = 'application/gzip' + resp.headers['X-Mask-Voxels'] = str(int(mask.sum())) + resp.headers['Cross-Origin-Resource-Policy'] = 'cross-origin' + return resp + except ValueError as ve: + return jsonify({"error": str(ve)}), 400 + except Exception as error: + print("[interactive_segment error]", type(error).__name__, error) + return jsonify({"error": "Interactive segmentation failed."}), 500 + finally: + _ANALYSIS_SLOTS.release() + + +@api_blueprint.route('/vessel-cpr/', methods=['POST']) +def vessel_cpr(case_id): + """Straightened vessel reformat + tumour-contact metrics for staging. + + Body JSON: { vessel_label:int, lesion_label?:int (default 22 pancreatic + lesion), res?, slab_radius_mm?, window? }. Returns metrics + + the reformat as a base64 PNG (already windowed for display). + """ + if not _ANALYSIS_SLOTS.acquire(blocking=False): + return jsonify(_ANALYSIS_BUSY_RESPONSE[0]), _ANALYSIS_BUSY_RESPONSE[1] + try: + import numpy as np + from services.advanced_analysis import analyze_vessel + body = request.get_json(force=True, silent=True) or {} + vessel_label = int(body.get("vessel_label", 0)) + if vessel_label <= 0: + return jsonify({"error": "vessel_label is required."}), 400 + lesion_label = int(body.get("lesion_label", 22)) + low = (body.get("res") or "low").lower() == "low" + + ct_path = _case_ct_path(case_id, low=low) + mask_path = _case_mask_path(case_id, low=low) + if not (os.path.exists(ct_path) and os.path.exists(mask_path)): + return jsonify({"error": "CT or segmentation not found for this case."}), 404 + + ct_obj = nib.load(ct_path) + # float32 (see interactive_segment): halves the per-request RAM footprint. + ct = ct_obj.get_fdata(dtype=np.float32) + labels = nib.load(mask_path).get_fdata(dtype=np.float32) + vessel_mask = (np.round(labels) == vessel_label).astype(np.uint8) + if vessel_mask.sum() == 0: + return jsonify({"error": "That vessel isn't segmented in this case."}), 422 + lesion_mask = (np.round(labels) == lesion_label).astype(np.uint8) + has_lesion = lesion_mask.sum() > 0 + + res = analyze_vessel( + ct, ct_obj.affine, vessel_mask, + lesion_mask if has_lesion else None, + slab_radius_mm=float(body.get("slab_radius_mm", 20.0)), + ) + + # Window the reformat (default soft-tissue) and PNG-encode it. + w = body.get("window") or {} + width = float(w.get("width", 400)); center = float(w.get("center", 40)) + lo, hi = center - width / 2, center + width / 2 + img = np.clip((res["reformat"] - lo) / max(hi - lo, 1e-6), 0, 1) + img8 = (img * 255).astype(np.uint8) + + from PIL import Image + png = io.BytesIO() + Image.fromarray(img8, mode="L").save(png, format="PNG") + import base64 + data_url = "data:image/png;base64," + base64.b64encode(png.getvalue()).decode() + + return jsonify({ + "length_mm": round(res["length_mm"], 1), + "max_contact_deg": round(res["max_contact_deg"], 1), + "contact_length_mm": round(res["contact_length_mm"], 1), + "has_lesion": bool(has_lesion), + "num_points": res["num_points"], + "contact_profile": [round(v, 1) for v in res["contact_profile"]], + "reformat_png": data_url, + "reformat_size": list(res["reformat"].shape), + }) + except ValueError as ve: + return jsonify({"error": str(ve)}), 400 + except Exception as error: + print("[vessel_cpr error]", type(error).__name__, error) + return jsonify({"error": "Vessel analysis failed."}), 500 + finally: + _ANALYSIS_SLOTS.release() diff --git a/flask-server/services/auto_segmentor.py b/flask-server/services/auto_segmentor.py index d32f34d..d389d98 100644 --- a/flask-server/services/auto_segmentor.py +++ b/flask-server/services/auto_segmentor.py @@ -1,5 +1,6 @@ import os import uuid +import signal import subprocess import re import csv @@ -14,6 +15,86 @@ # Only one model inference runs at a time to avoid GPU OOM _gpu_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 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.wait(timeout=5) + except Exception: + 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: try: @@ -52,13 +133,36 @@ 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, precision="fp16"): """ 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) + # Precision applies to the GPU inference stages only. The segmentation models + # (ePAI, SuPreM, MedFormer, R-Super, Atlas-Net) and the OpenVAE preprocessor are + # torch networks that can run in FP16; ShapeKit is CPU-only shape refinement, so + # precision is not meaningful there and is pinned to fp32. The value is exported so + # the inference subprocess (which inherits this environment) can load the matching + # half-precision weights. + effective_precision = "fp32" if model == "ShapeKit" else (precision or "fp16") + os.environ["BODYMAPS_PRECISION"] = effective_precision + print(f"[INFO] Inference precision for {model}: {effective_precision}") if model == 'ePAI': conda_path = _resolve_conda_activate_path() return _run_epai_inference( @@ -363,7 +467,7 @@ 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: - subprocess.run( + _tracked_run( full_cmd, shell=True, executable="/bin/bash", @@ -429,15 +533,17 @@ def _run_suprem_inference(input_path: str, session_dir: str) -> str: f"python -W ignore {shlex.quote(os.path.join(suprem_src, 'inference.py'))} " f"--data_root_path {shlex.quote(inputs_dir)} " f"--save_dir {shlex.quote(output_dir)} " - f"--resume {shlex.quote(checkpoint)} " + f"--checkpoint {shlex.quote(checkpoint)} " + f"--backbone unet " + f"--suprem " f"--store_result" ) print(f"[INFO] Running SuPreM native inference") print(full_cmd) try: - subprocess.run(full_cmd, shell=True, executable="/bin/bash", check=True, - 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}" @@ -496,7 +602,7 @@ def _run_openvae_inference(input_path: str, session_dir: str) -> str: ) print(f"[INFO] Running OpenVAE inference\n{full_cmd}") try: - subprocess.run(full_cmd, shell=True, executable="/bin/bash", check=True, 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}" @@ -588,7 +694,7 @@ def _run_medformer_inference(input_path: str, session_dir: str) -> str: ) print(f"[INFO] Running MedFormer inference\n{full_cmd}") try: - subprocess.run( + _tracked_run( full_cmd, shell=True, executable="/bin/bash", check=True, cwd=rsuper_src, stdout=subprocess.PIPE, stderr=subprocess.STDOUT ) @@ -651,7 +757,7 @@ def _run_rsuper_inference(input_path: str, session_dir: str) -> str: ) print(f"[INFO] Running R-Super inference\n{full_cmd}") try: - subprocess.run( + _tracked_run( full_cmd, shell=True, executable="/bin/bash", check=True, cwd=rsuper_src, stdout=subprocess.PIPE, stderr=subprocess.STDOUT ) @@ -714,7 +820,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: - subprocess.run(full_cmd, shell=True, executable="/bin/bash", check=True) + _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}" @@ -738,7 +844,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}" @@ -953,7 +1059,7 @@ def _run_shapekit_inference(input_dir: str, session_dir: str) -> str: ) print(f"[INFO] Running ShapeKit post-processing\n{full_cmd}") try: - subprocess.run(full_cmd, shell=True, executable="/bin/bash", check=True, 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}"