cvpilot-tool/apps/desktop/electron/ipcEvent/alg/index.ts

import { downloadSeriesDicomFiles, getTotalScanLength } from "../../core/pacs";
import { executeInferTask } from "../../core/alg";
import { InferDeviceEnum, InferStructuralEnum } from "../../core/alg.type";
import { db } from "../../core/db";
import log from "electron-log";
import path from "node:path";
import { app, ipcMain } from "electron";
import { findSTLFiles } from "./util";
import { readFileSync } from "node:fs";

export const registerAlgHandler = () => {
  ipcMain.handle("device:infer:set", async (_event, inferDevice) => {
    try {
      await db.update(({ setting }) => ({ ...setting, inferDevice }));
      return { success: true, msg: `推理硬件修改为${inferDevice}` };
    } catch (error) {
      await db.update(({ setting }) => ({ ...setting, inferDevice }));
      log.error(error);
      return { success: false, msg: `操作失败` };
    }
  });

  ipcMain.on("model:infer", async (event, SeriesInstanceUIDs) => {
    // 构造推理任务参数列表
    const pu = InferDeviceEnum.GPU;

    const turbo = false;
    const seg_schedule = true;
    for (let i = 0; i < SeriesInstanceUIDs.length; i++) {
      const SeriesInstanceUID = SeriesInstanceUIDs[i];
      const physicalLength = await getTotalScanLength(SeriesInstanceUID);
      const module =
        physicalLength && physicalLength < 200
          ? InferStructuralEnum.AORTA
          : InferStructuralEnum.PERI;
      // 下载dicom到本地，获取文件夹路径
      const img_path = await downloadSeriesDicomFiles(SeriesInstanceUID);
      const save_path = path.join(
        app.getPath("userData"),
        "output",
        SeriesInstanceUID
      );
      const task = { save_path, pu, module, turbo, seg_schedule, img_path };
      console.log("infer task: ", task);
      await executeInferTask(task, (data) => {
        console.log(data);
        event.reply("infer:progress", data);
      });
    }
  });

  ipcMain.handle("alg:assets", async (_event, SeriesInstanceUID) => {
    const physicalLength = await getTotalScanLength(SeriesInstanceUID);
    const module =
      physicalLength && physicalLength < 200
        ? InferStructuralEnum.AORTA
        : InferStructuralEnum.PERI;
    const rootPath = path.join(
      app.getPath("userData"),
      "output",
      SeriesInstanceUID
    );
    const stlsPath = path.join(rootPath, module, "visualization");
    const stls = await findSTLFiles(stlsPath);
    if (stls.length > 0) {
      try {
        let result = {}
        // 读取 STL 文件并转换为 ArrayBuffer
        const stlFiles = stls.map((file) => {
          const filePath = path.join(stlsPath, file);
          const data = readFileSync(filePath);
          return {
            fileName: file,
            data: data.buffer,
          };
        });
        if (module === InferStructuralEnum.AORTA) {
          // 读取测量json
          const measurementPath = path.join(rootPath, module, "measurement.json");
          const measurementData = readFileSync(measurementPath, "utf-8");
          result = { stlFiles, measurement: JSON.parse(measurementData) }
        } else if (module === InferStructuralEnum.PERI) {
          result = { stlFiles }
        }
        return { ...result, module };
      } catch (error) {
        console.error("Error reading measurement.json:", error);
        throw new Error("Failed to read or parse measurement.json");
      }
    }
  });
};