"""Helpers for seeding the bundled demo media project.""" from __future__ import annotations import os import shutil import tempfile from pathlib import Path import cv2 from sqlalchemy.orm import Session from minio_client import upload_file from models import Frame, Project, User from services.frame_parser import natural_filename_key, parse_dicom, upload_frames_to_minio from statuses import PROJECT_STATUS_PENDING, PROJECT_STATUS_READY DEMO_DICOM_PROJECT_NAME = "演视DICOM序列" DEMO_DICOM_PARSE_FPS = 30.0 DEMO_VIDEO_PROJECT_NAME = "演视LC视频序列" LEGACY_DEMO_VIDEO_PROJECT_NAMES = {"Data_MyVideo_1"} LEGACY_DEMO_DICOM_PROJECT_NAMES = {"演示DICOM序列"} def demo_dicom_files(dicom_dir: str) -> list[Path]: """Return .dcm files in natural file-name order.""" root = Path(dicom_dir) if not root.exists() or not root.is_dir(): return [] return sorted( [path for path in root.iterdir() if path.is_file() and path.name.lower().endswith(".dcm")], key=lambda path: natural_filename_key(path.name), ) def create_unparsed_video_demo_project( db: Session, *, owner: User, video_path: str, project_name: str = DEMO_VIDEO_PROJECT_NAME, ) -> Project: """Create the bundled demo video project without extracting frames.""" source = Path(video_path) if not source.exists() or not source.is_file(): raise FileNotFoundError(f"Demo video not found: {video_path}") project = Project( name=project_name, description="默认演示视频,尚未生成帧", status=PROJECT_STATUS_PENDING, source_type="video", parse_fps=30.0, original_fps=None, owner_user_id=owner.id, ) db.add(project) db.flush() data = source.read_bytes() object_name = f"uploads/{project.id}/{source.name}" upload_file(object_name, data, content_type="video/mp4", length=len(data)) project.video_path = object_name project.thumbnail_url = None db.commit() db.refresh(project) return project def create_parsed_dicom_demo_project( db: Session, *, owner: User, dicom_dir: str, project_name: str = DEMO_DICOM_PROJECT_NAME, ) -> Project: """Create the demo DICOM project, upload the series, and register parsed frames.""" dcm_files = demo_dicom_files(dicom_dir) if not dcm_files: raise FileNotFoundError(f"Demo DICOM series not found: {dicom_dir}") project = Project( name=project_name, description=f"默认演示 DICOM 序列,已按文件名自然顺序生成 {len(dcm_files)} 帧", status=PROJECT_STATUS_PENDING, source_type="dicom", parse_fps=DEMO_DICOM_PARSE_FPS, original_fps=None, owner_user_id=owner.id, ) db.add(project) db.flush() dicom_prefix = f"uploads/{project.id}/dicom" for dcm_file in dcm_files: data = dcm_file.read_bytes() upload_file( f"{dicom_prefix}/{dcm_file.name}", data, content_type="application/dicom", length=len(data), ) project.video_path = dicom_prefix tmp_dir = tempfile.mkdtemp(prefix=f"seg_demo_dicom_{project.id}_") try: output_dir = os.path.join(tmp_dir, "frames") frame_files = parse_dicom(dicom_dir, output_dir) object_names = upload_frames_to_minio(frame_files, project.id) for idx, obj_name in enumerate(object_names): image = cv2.imread(frame_files[idx]) height, width = image.shape[:2] if image is not None else (None, None) db.add(Frame( project_id=project.id, frame_index=idx, image_url=obj_name, width=width, height=height, timestamp_ms=idx * 1000.0 / DEMO_DICOM_PARSE_FPS, source_frame_number=idx, )) if object_names: project.thumbnail_url = object_names[0] project.status = PROJECT_STATUS_READY db.commit() db.refresh(project) return project finally: shutil.rmtree(tmp_dir, ignore_errors=True)