import os
import pydicom
import numpy as np
import cv2
import matplotlib.pyplot as plt
import matplotlib as mpl
from matplotlib.colors import LinearSegmentedColormap
from scipy import ndimage
import glob
import warnings
from natsort import natsorted
import time
import traceback
warnings.filterwarnings("ignore")
# 创建并注册NIH颜色映射
nih_colors = [
(0.0, 0.0, 0.0),
(0.0, 0.0, 1.0),
(0.0, 1.0, 1.0),
(0.486, 0.988, 0.0),
(1.0, 1.0, 0.0),
(1.0, 0.0, 0.0),
]
nih_cmap = LinearSegmentedColormap.from_list("nih", nih_colors)
# 注册颜色映射到matplotlib
try:
mpl.colormaps.register(cmap=nih_cmap)
except AttributeError:
try:
from matplotlib import cm
cm.register_cmap(name="nih", cmap=nih_cmap)
except:
pass
def load_dicom_slice(file_path):
"""加载单个DICOM切片"""
try:
ds = pydicom.dcmread(file_path)
img = ds.pixel_array.astype(np.float32)
slope = float(getattr(ds, "RescaleSlope", 1))
intercept = float(getattr(ds, "RescaleIntercept", 0))
return img * slope + intercept
except Exception as e:
print(f"错误 加载切片 {file_path} 失败: {str(e)}")
return None
def load_dicom_volume(dicom_folder_path):
"""加载DICOM文件夹中的三维图像数据(串行版本)"""
print(f"正在加载 {dicom_folder_path}")
# 获取所有DICOM文件
dicom_files = glob.glob(os.path.join(dicom_folder_path, "*.IMA")) + \
glob.glob(os.path.join(dicom_folder_path, "*.dcm"))
if not dicom_files:
print(f"警告 在 {dicom_folder_path} 中未找到DICOM文件")
return None, None
# 使用自然排序
dicom_files = natsorted(dicom_files)
# 获取Z位置并排序
z_positions = []
valid_files = []
for f in dicom_files:
try:
ds = pydicom.dcmread(f, stop_before_pixels=True)
z_pos = float(ds.ImagePositionPatient[2])
z_positions.append(z_pos)
valid_files.append(f)
except Exception as e:
print(f"警告 读取 {f} 的元数据失败: {str(e)}")
if not z_positions:
print(f"错误 无法获取任何DICOM文件的Z位置")
return None, None
# 按Z位置排序文件
sorted_indices = np.argsort(z_positions)
dicom_files = [valid_files[i] for i in sorted_indices]
z_positions = [z_positions[i] for i in sorted_indices]
# 加载第一个切片获取图像尺寸
first_slice = load_dicom_slice(dicom_files[0])
if first_slice is None:
print("错误 无法加载第一个DICOM切片")
return None, None
image_shape = first_slice.shape
num_slices = len(dicom_files)
# 预分配内存
image_3d = np.zeros((image_shape[0], image_shape[1], num_slices), dtype=np.float32)
# 串行加载所有切片
for i, file_path in enumerate(dicom_files):
print(f"加载切片 {i+1}/{num_slices}", end="\r")
slice_data = load_dicom_slice(file_path)
if slice_data is not None:
image_3d[:, :, i] = slice_data
print(f"\n成功加载 {num_slices} 个切片")
# 计算间距
z_spacing = np.mean(np.diff(z_positions)) if len(z_positions) > 1 else 1.0
try:
ds = pydicom.dcmread(dicom_files[0])
pixel_spacing = [float(sp) for sp in ds.PixelSpacing]
except:
pixel_spacing = [1.0, 1.0]
spacing = (z_spacing, pixel_spacing[0], pixel_spacing[1]) # (Z, Y, X)
print(f"图像尺寸 (Y, X, Z) {image_3d.shape}")
print(f"体素空间间距 (Z, Y, X) {spacing} 单位 mm")
return image_3d, spacing
def render_projection(data, angle):
"""优化渲染函数 - 使用2D旋转和最大强度投影"""
# 沿Z轴进行最大强度投影
projection = np.max(data, axis=2)
# 旋转投影图像
rotated = ndimage.rotate(projection, angle, reshape=False, order=1, mode='constant', cval=0.0)
return rotated
def generate_rotation_frames(
data, output_folder, prefix, angle_step=10, intensity_limit=None
):
print(f"正在生成旋转帧 保存到 {output_folder}")
os.makedirs(output_folder, exist_ok=True)
# 计算强度限制
if intensity_limit is None:
non_zero_data = data[data > 0]
if len(non_zero_data) > 0:
intensity_limit = np.percentile(non_zero_data, 95)
else:
intensity_limit = np.max(data) * 0.8 if np.max(data) > 0 else 1
print(f"使用强度限制: {intensity_limit:.2f}")
angles = list(range(0, 360, angle_step))
frame_files = []
generated_count = 0
skipped_count = 0
for angle in angles:
filename = os.path.join(output_folder, f"{prefix}_{angle:03d}.png")
frame_files.append(filename)
if os.path.exists(filename):
print(f"跳过已存在的帧 {angle}°")
skipped_count += 1
continue
start_time = time.time()
# 渲染投影
rotated_data = render_projection(data, angle)
# 应用强度限制
if intensity_limit > 0:
rotated_data = np.clip(rotated_data, 0, intensity_limit)
# 保存图像
plt.imsave(filename, rotated_data, cmap="nih", dpi=150, origin='lower')
elapsed = time.time() - start_time
print(f"已生成帧 {angle}° - 耗时: {elapsed:.2f}秒")
generated_count += 1
print(f"总计 跳过 {skipped_count} 帧 生成 {generated_count} 帧")
return frame_files
def create_video_from_frames(frame_files, output_video, frame_rate=15):
print(f"正在创建视频 {output_video}")
if not frame_files:
print("错误: 没有找到帧图像文件")
return
# 使用OpenCV直接读取图像并创建视频
first_frame = cv2.imread(frame_files[0])
if first_frame is None:
print(f"错误: 无法读取第一帧: {frame_files[0]}")
return
height, width = first_frame.shape[:2]
# 尝试多种视频编码
codecs = ['mp4v', 'avc1', 'X264']
video_writer = None
for codec in codecs:
try:
fourcc = cv2.VideoWriter_fourcc(*codec)
video_writer = cv2.VideoWriter(output_video, fourcc, frame_rate, (width, height))
if video_writer.isOpened():
print(f"使用编解码器: {codec}")
break
except:
pass
if video_writer is None or not video_writer.isOpened():
print("错误: 无法创建视频文件,没有可用的编解码器")
return
for i, frame_path in enumerate(frame_files):
frame = cv2.imread(frame_path)
if frame is None:
print(f"警告: 无法读取帧 {frame_path}, 跳过")
continue
video_writer.write(frame)
if (i + 1) % 10 == 0 or (i + 1) == len(frame_files):
print(f"已处理 {i+1}/{len(frame_files)} 帧")
video_writer.release()
print(f"视频已保存 {output_video}")
# 清理临时帧文件
for frame_file in frame_files:
try:
os.remove(frame_file)
except:
pass
def generate_full_dose_video():
print("生成全剂量患者的三维影像可旋转视频")
base_path = "LAFOV-Static-PET"
full_dose_path = os.path.join(base_path, "Full_dose")
if not os.path.exists(full_dose_path):
print(f"错误 全剂量数据路径不存在 - {full_dose_path}")
return
volume, spacing = load_dicom_volume(full_dose_path)
if volume is None:
print("错误 无法加载全剂量数据")
return
output_folder = "full_dose_frames"
frame_files = generate_rotation_frames(
volume, output_folder, "full_dose", angle_step=10
)
if not frame_files:
print("错误: 没有生成任何帧")
return
output_video = "full_dose_rotation_video.mp4"
create_video_from_frames(frame_files, output_video)
print(f"全剂量旋转视频已完成 {output_video}")
def generate_multi_dose_comparison_video():
print("生成不同浓度患者影像的并排可旋转视频")
base_path = "LAFOV-Static-PET"
# 使用原始文件夹名称(包含空格)
dose_folders = {
"全剂量": "Full_dose",
"1/2剂量": "1-2 dose",
"1/4剂量": "1-4 dose",
"1/10剂量": "1-10 dose",
"1/20剂量": "1-20 dose",
"1/50剂量": "1-50 dose",
"1/100剂量": "1-100 dose",
}
volumes = {}
for dose_name, folder_name in dose_folders.items():
folder_path = os.path.join(base_path, folder_name)
if os.path.exists(folder_path):
print(f"正在加载 {dose_name} 数据...")
volume, spacing = load_dicom_volume(folder_path)
if volume is not None:
volumes[dose_name] = volume
print(f"已加载 {dose_name}")
else:
print(f"跳过 {dose_name} 无法加载数据")
else:
print(f"跳过 {dose_name} 路径不存在: {folder_path}")
if not volumes:
print("错误 未找到任何有效的数据")
return
# 使用全剂量数据计算强度限制
if "全剂量" in volumes:
full_dose_volume = volumes["全剂量"]
non_zero_data = full_dose_volume[full_dose_volume > 0]
if len(non_zero_data) > 0:
intensity_limit = np.percentile(non_zero_data, 95)
else:
intensity_limit = np.max(full_dose_volume) * 0.8
else:
first_volume = next(iter(volumes.values()))
non_zero_data = first_volume[first_volume > 0]
if len(non_zero_data) > 0:
intensity_limit = np.percentile(non_zero_data, 95)
else:
intensity_limit = np.max(first_volume) * 0.8
print(f"使用强度限制: {intensity_limit:.2f}")
output_folder = "multi_dose_comparison_frames"
os.makedirs(output_folder, exist_ok=True)
angles = list(range(0, 360, 10))
frame_files = []
generated_count = 0
skipped_count = 0
# 预渲染所有剂量和角度的投影
precomputed_projections = {}
for dose_name, volume in volumes.items():
projections = []
print(f"预渲染 {dose_name} 投影...")
for angle in angles:
projections.append(render_projection(volume, angle))
precomputed_projections[dose_name] = projections
# 生成并排比较图像
for i, angle in enumerate(angles):
filename = os.path.join(output_folder, f"comparison_{angle:03d}.png")
frame_files.append(filename)
if os.path.exists(filename):
print(f"跳过已存在的并排帧 {angle}°")
skipped_count += 1
continue
start_time = time.time()
num_doses = len(volumes)
fig, axes = plt.subplots(1, num_doses, figsize=(4 * num_doses, 4))
if num_doses == 1:
axes = [axes]
for j, dose_name in enumerate(volumes.keys()):
img = precomputed_projections[dose_name][i]
im = axes[j].imshow(img, cmap="nih", vmin=0, vmax=intensity_limit, origin='lower')
axes[j].set_title(f"{dose_name}", fontsize=10)
axes[j].axis("off")
plt.colorbar(im, ax=axes[j], fraction=0.046, pad=0.04)
plt.tight_layout(pad=1.0)
plt.savefig(filename, dpi=150, bbox_inches='tight')
plt.close(fig)
elapsed = time.time() - start_time
print(f"已生成并排帧 {angle}° - 耗时: {elapsed:.2f}秒")
generated_count += 1
print(f"总计 跳过 {skipped_count} 帧 生成 {generated_count} 帧")
if not frame_files:
print("错误: 没有生成任何帧")
return
output_video = "multi_dose_comparison_video.mp4"
create_video_from_frames(frame_files, output_video)
print(f"多剂量并排旋转视频已完成 {output_video}")
def main():
print("PET图像三维可旋转视频生成")
print("=" * 50)
plt.rcParams["font.sans-serif"] = ["SimHei", "DejaVu Sans"]
plt.rcParams["axes.unicode_minus"] = False
try:
start_time = time.time()
generate_full_dose_video()
print("\n" + "=" * 50)
print(f"全剂量视频生成耗时: {time.time() - start_time:.2f}秒")
start_time = time.time()
generate_multi_dose_comparison_video()
print("\n" + "=" * 50)
print(f"多剂量对比视频生成耗时: {time.time() - start_time:.2f}秒")
print("所有视频生成完成")
except Exception as e:
print(f"错误 {str(e)}")
traceback.print_exc()
if __name__ == "__main__":
main()
最新发布