2026-05-20-02-32-47 支持NII导出包与分割类别范围

2026-05-20 02:40:50 +08:00
parent 66ad99f996
commit 68fb0cb564
7 changed files with 469 additions and 29 deletions
--- a/WebSite/server.ts
+++ b/WebSite/server.ts
@@ -9,6 +9,8 @@ import { fileURLToPath } from 'node:url';
 type ProjectStatus = 'pending' | 'completed' | 'processing';
 type DicomPlane = 'axial' | 'sagittal' | 'coronal';
 type DicomDisplayMode = 'default' | 'bone' | 'soft' | 'contrast';
+type ProjectExportTarget = 'dicom' | 'segmentation' | 'pose';
+type SegmentationExportScope = 'all' | 'visible';

 interface ModuleStyleRecord {
  visible: boolean;
@@ -772,7 +774,20 @@ function fillExportRows(data: Buffer, width: number, height: number, slice: numb
  });
 }

-function createSegmentationData(project: ProjectRecord, volume: DicomHuVolume, pose: ModelPoseValue) {
+function getModuleStyle(project: ProjectRecord, fileName: string, index: number): ModuleStyleRecord {
+  return project.moduleStyles[fileName] ?? {
+    visible: true,
+    color: defaultModuleColors[index % defaultModuleColors.length],
+    opacity: 0.72,
+    partId: index + 1,
+  };
+}
+
+function isModuleIncludedForExport(style: ModuleStyleRecord, scope: SegmentationExportScope) {
+  return scope === 'all' || style.visible !== false;
+}
+
+function createSegmentationData(project: ProjectRecord, volume: DicomHuVolume, pose: ModelPoseValue, scope: SegmentationExportScope = 'visible') {
  const data = Buffer.alloc(volume.width * volume.height * volume.depth);
  const previews = (project.stlFiles ?? []).reduce<Record<string, ModelPreviewRecord>>((accumulator, fileName) => {
    const filePath = path.join(modelDir, fileName);
@@ -803,14 +818,9 @@ function createSegmentationData(project: ProjectRecord, volume: DicomHuVolume, p

  (project.stlFiles ?? []).forEach((fileName, index) => {
    const payload = previews[fileName];
-    const style = project.moduleStyles[fileName] ?? {
-      visible: true,
-      color: defaultModuleColors[index % defaultModuleColors.length],
-      opacity: 0.72,
-      partId: index + 1,
-    };
+    const style = getModuleStyle(project, fileName, index);

-    if (!payload || style.visible === false) {
+    if (!payload || !isModuleIncludedForExport(style, scope)) {
      return;
    }

@@ -872,6 +882,55 @@ function createSegmentationData(project: ProjectRecord, volume: DicomHuVolume, p
  return data;
 }

+function createSegmentationLabelMetadata(project: ProjectRecord, scope: SegmentationExportScope, activePose?: ModelPoseValue) {
+  const labels = (project.stlFiles ?? [])
+    .map((fileName, index) => {
+      const style = getModuleStyle(project, fileName, index);
+      if (!isModuleIncludedForExport(style, scope)) {
+        return null;
+      }
+      const name = fileName.replace(/\.stl$/i, '');
+      const label = clampNumber(Math.round(style.partId || index + 1), 1, 255);
+
+      return {
+        label,
+        partId: label,
+        name,
+        categoryName: name,
+        className: name,
+        fileName,
+        color: style.color,
+        opacity: style.opacity,
+        visible: style.visible !== false,
+      };
+    })
+    .filter((item): item is {
+      label: number;
+      partId: number;
+      name: string;
+      categoryName: string;
+      className: string;
+      fileName: string;
+      color: string;
+      opacity: number;
+      visible: boolean;
+    } => Boolean(item));
+
+  return Buffer.from(JSON.stringify({
+    project: {
+      id: project.id,
+      name: project.name,
+      dicomPath: project.dicomPath,
+      modelPath: project.modelPath,
+    },
+    generatedAt: now(),
+    segmentationScope: scope,
+    activePose: activePose ?? null,
+    labels,
+    note: 'Label values correspond to ReVoxelSeg STL component hierarchy partId values.',
+  }, null, 2), 'utf8');
+}
+
 function parseModelPoseQuery(raw: unknown) {
  if (typeof raw !== 'string' || !raw.trim()) {
    return undefined;
@@ -884,13 +943,37 @@ function parseModelPoseQuery(raw: unknown) {
  }
 }

-function createNiftiExport(project: ProjectRecord, files: string[], target: 'dicom' | 'segmentation', compressed: boolean, pose?: ModelPoseValue) {
+function parseSegmentationScope(raw: unknown): SegmentationExportScope {
+  return raw === 'all' ? 'all' : 'visible';
+}
+
+function parseExportTargets(raw: unknown): ProjectExportTarget[] {
+  const values = typeof raw === 'string' ? raw.split(',') : [];
+  const targets = values.filter((value): value is ProjectExportTarget => (
+    value === 'dicom' || value === 'segmentation' || value === 'pose'
+  ));
+  return [...new Set(targets)];
+}
+
+function createNiftiExport(
+  project: ProjectRecord,
+  files: string[],
+  target: 'dicom' | 'segmentation',
+  compressed: boolean,
+  pose?: ModelPoseValue,
+  segmentationScope: SegmentationExportScope = 'visible',
+) {
  const volume = readDicomHuVolume(files);
  if (target === 'dicom') {
    return createNiftiBuffer(volume, volume.data, 'dicom', compressed);
  }

-  return createNiftiBuffer(volume, createSegmentationData(project, volume, pose ?? defaultModelPose), 'segmentation', compressed);
+  return createNiftiBuffer(
+    volume,
+    createSegmentationData(project, volume, pose ?? defaultModelPose, segmentationScope),
+    'segmentation',
+    compressed,
+  );
 }

 function createPoseExport(project: ProjectRecord, activePose?: ModelPoseValue) {
@@ -909,6 +992,64 @@ function createPoseExport(project: ProjectRecord, activePose?: ModelPoseValue) {
  }, null, 2), 'utf8');
 }

+function createProjectExportBundle({
+  project,
+  files,
+  targets,
+  compressed,
+  activePose,
+  segmentationScope,
+}: {
+  project: ProjectRecord;
+  files: string[];
+  targets: ProjectExportTarget[];
+  compressed: boolean;
+  activePose?: ModelPoseValue;
+  segmentationScope: SegmentationExportScope;
+}) {
+  const entries: Array<{ name: string; data: Buffer; mtime?: number }> = [];
+  const needsVolume = targets.includes('dicom') || targets.includes('segmentation');
+  const volume = needsVolume ? readDicomHuVolume(files) : null;
+  const format = compressed ? 'nii.gz' : 'nii';
+  const exportRoot = `${project.id}-nifti-export`;
+
+  if (targets.includes('dicom') && volume) {
+    entries.push({
+      name: `${exportRoot}/${project.id}-dicom-image.${format}`,
+      data: createNiftiBuffer(volume, volume.data, 'dicom', compressed),
+    });
+  }
+
+  if (targets.includes('segmentation') && volume) {
+    entries.push({
+      name: `${exportRoot}/${project.id}-segmentation-label.${format}`,
+      data: createNiftiBuffer(
+        volume,
+        createSegmentationData(project, volume, activePose ?? defaultModelPose, segmentationScope),
+        'segmentation',
+        compressed,
+      ),
+    });
+    entries.push({
+      name: `${exportRoot}/${project.id}-segmentation-labels.json`,
+      data: createSegmentationLabelMetadata(project, segmentationScope, activePose),
+    });
+  }
+
+  if (targets.includes('pose')) {
+    entries.push({
+      name: `${exportRoot}/${project.id}-pose-data.json`,
+      data: createPoseExport(project, activePose),
+    });
+  }
+
+  if (!entries.length) {
+    throw new Error('未选择可导出的内容');
+  }
+
+  return createTarGz(entries);
+}
+
 function findProject(state: AppState, projectId: string) {
  return state.projects.find((candidate) => candidate.id === projectId);
 }
@@ -1443,14 +1584,12 @@ function createTarEntryHeader(name: string, size: number, mtime: number) {
  return header;
 }

-function createDicomTarGz(files: string[]) {
+function createTarGz(entries: Array<{ name: string; data: Buffer; mtime?: number }>) {
  const chunks: Buffer[] = [];

-  files.forEach((fileName) => {
-    const filePath = path.join(dicomDir, fileName);
-    const stat = fs.statSync(filePath);
-    const data = fs.readFileSync(filePath);
-    chunks.push(createTarEntryHeader(`Head_CT_DICOM/${fileName}`, data.length, stat.mtimeMs / 1000));
+  entries.forEach((entry) => {
+    const data = entry.data;
+    chunks.push(createTarEntryHeader(entry.name, data.length, entry.mtime ?? Date.now() / 1000));
    chunks.push(data);
    const remainder = data.length % 512;
    if (remainder > 0) {
@@ -1462,6 +1601,18 @@ function createDicomTarGz(files: string[]) {
  return zlib.gzipSync(Buffer.concat(chunks));
 }

+function createDicomTarGz(files: string[]) {
+  return createTarGz(files.map((fileName) => {
+    const filePath = path.join(dicomDir, fileName);
+    const stat = fs.statSync(filePath);
+    return {
+      name: `Head_CT_DICOM/${fileName}`,
+      data: fs.readFileSync(filePath),
+      mtime: stat.mtimeMs / 1000,
+    };
+  }));
+}
+
 function estimateSliceSpacingFromAttributes(attributes: DicomAttributes[]) {
  const diffs: number[] = [];
  for (let index = 1; index < attributes.length; index += 1) {
@@ -1889,6 +2040,7 @@ async function startServer() {
    const requestedTarget = targetOverride ?? String(req.query.target ?? 'segmentation');
    const target = requestedTarget === 'dicom' || requestedTarget === 'pose' ? requestedTarget : 'segmentation';
    const activePose = parseModelPoseQuery(req.query.pose);
+    const segmentationScope = parseSegmentationScope(req.query.segmentationScope);

    try {
      if (target === 'pose') {
@@ -1904,7 +2056,7 @@ async function startServer() {
      const files = getProjectDicomFiles(project);
      const format = req.query.format === 'nii' ? 'nii' : 'nii.gz';
      const compressed = format === 'nii.gz';
-      const payload = createNiftiExport(project, files, target, compressed, activePose);
+      const payload = createNiftiExport(project, files, target, compressed, activePose, segmentationScope);
      const suffix = target === 'dicom' ? 'dicom-image' : 'segmentation-label';
      const filename = `${project.id}-${suffix}.${format}`;
      fs.writeFileSync(path.join(exportDir, filename), payload);
@@ -1919,6 +2071,49 @@ async function startServer() {
    }
  };

+  app.get('/api/projects/:projectId/export-bundle', (req, res) => {
+    const state = readState();
+    const project = state.projects.find((candidate) => candidate.id === req.params.projectId);
+
+    if (!project) {
+      res.status(404).json({ message: '项目不存在' });
+      return;
+    }
+
+    const targets = parseExportTargets(req.query.targets);
+    if (!targets.length) {
+      res.status(400).json({ message: '请至少选择一个导出内容' });
+      return;
+    }
+
+    const activePose = parseModelPoseQuery(req.query.pose);
+    const segmentationScope = parseSegmentationScope(req.query.segmentationScope);
+    const format = req.query.format === 'nii' ? 'nii' : 'nii.gz';
+    const compressed = format === 'nii.gz';
+
+    try {
+      const files = getProjectDicomFiles(project);
+      const payload = createProjectExportBundle({
+        project,
+        files,
+        targets,
+        compressed,
+        activePose,
+        segmentationScope,
+      });
+      const filename = `${project.id}-nifti-export.tar.gz`;
+      fs.writeFileSync(path.join(exportDir, filename), payload);
+      project.exportedMaskCount += targets.includes('segmentation') ? 1 : 0;
+      writeState(state);
+
+      res.setHeader('Content-Type', 'application/gzip');
+      res.setHeader('Content-Disposition', `attachment; filename="${filename}"`);
+      res.send(payload);
+    } catch (error) {
+      res.status(422).json({ message: error instanceof Error ? error.message : '导出包生成失败' });
+    }
+  });
+
  app.get('/api/projects/:projectId/export-nifti', (req, res) => handleProjectExport(req, res));
  app.get('/api/projects/:projectId/export-mask', (req, res) => handleProjectExport(req, res, 'segmentation'));
  app.post('/api/projects/:projectId/export-mask', (req, res) => handleProjectExport(req, res, 'segmentation'));