|
- import os.path as osp
-
- import SimpleITK as sitk
- import cv2
- from cv2 import VideoWriter_fourcc
- import numpy as np
- from tqdm import tqdm
-
-
- def load_medical_data(f):
- """
- load data of different format into numpy array, return data is in xyz
- f: the complete path to the file that you want to load
- """
- filename = osp.basename(f).lower()
-
- if filename.endswith((".nii", ".nii.gz", ".dcm")):
- itkimage = sitk.ReadImage(f)
- if itkimage.GetDimension() == 4:
- slicer = sitk.ExtractImageFilter()
- s = list(itkimage.GetSize())
- s[-1] = 0
- slicer.SetSize(s)
- images = []
- for slice_idx in range(itkimage.GetSize()[-1]):
- slicer.SetIndex([0, 0, 0, slice_idx])
- sitk_volume = slicer.Execute(itkimage)
- images.append(sitk_volume)
- images = [sitk.DICOMOrient(img, "SLP") for img in images]
- f_nps = [sitk.GetArrayFromImage(img) for img in images]
- else:
- image = sitk.DICOMOrient(itkimage, "SLP")
- f_np = sitk.GetArrayFromImage(image)
- f_nps = [f_np]
-
- elif filename.endswith((".mha", ".mhd", "nrrd")):
- itkimage = sitk.DICOMOrient(sitk.ReadImage(f), "SLP")
- f_np = sitk.GetArrayFromImage(itkimage)
- if f_np.ndim == 4:
- f_nps = [f_np[:, :, :, idx] for idx in range(f_np.shape[3])]
- else:
- f_nps = [f_np]
- elif filename.endswith(".raw"):
- raise RuntimeError(
- f"Received {f}. Please only provide path to .mhd file, not to .raw file"
- )
- else:
- raise NotImplementedError
-
- return f_nps
-
-
- def normalize(frame, ww=400, wc=0):
- wl = wc - ww / 2
- wh = wc + ww / 2
- frame = frame.astype("float16")
- np.clip(frame, wl, wh, out=frame)
- frame = (frame - wl) / ww * 255
- frame = frame.astype(np.uint8)
- frame = cv2.cvtColor(frame, cv2.COLOR_GRAY2BGR)
- return frame
-
-
- def array_to_video(array_data, video_path, fps=15):
- h, w, s = array_data.shape
- fourcc = VideoWriter_fourcc(*"mp4v")
- videoWriter = cv2.VideoWriter(video_path, fourcc, fps, (w, h))
- for idx in tqdm(range(s)):
- frame = array_data[:, :, idx]
- frame = normalize(frame)
- videoWriter.write(frame)
-
- videoWriter.release()
- cv2.destroyAllWindows()
- print("转视频结束!")
-
-
- if __name__ == "__main__":
- path = "/home/lin/Desktop/data/coronacases_org_001.nii.gz"
- video_path = "/home/lin/Desktop/temp.mp4"
- total_data = load_medical_data(path)
- for data in total_data:
- array_to_video(data, video_path)
|