Skip to content

Instantly share code, notes, and snippets.

@jphdotam
Created April 3, 2019 11:30
Show Gist options
  • Save jphdotam/26d8b5799f7ed89669dc0305a7213312 to your computer and use it in GitHub Desktop.
Save jphdotam/26d8b5799f7ed89669dc0305a7213312 to your computer and use it in GitHub Desktop.
import os
import cv2
import glob
import pydicom
from pydicom.dicomdir import DicomDir
from pydicom.errors import InvalidDicomError
## EDIT THESE TWO LINES IF NEEDED:
FOLDER_PATH = "./dicom_study/" # Path to your folder of DICOMS; default assumes they're in a folder called 'dicom_study' in the same folder as this file
BLANK_PROPORTION = 0.05 # Set this to 0.1 to blank out top 10% of video (and therefore anonymise)
## You shouldn't need to edit anything below this
CONVERT_TO_BGR = True # This reverses the red and blue channels, because CV2 (god knows why) uses BGR not RGB
IGNORED_EXTENSIONS = [".PDF", ".XML", ".ANT", ".AVI", ".avi", ".jpg", ".JPG", '.mp4', '.MP4'] # Don't convert these files
file_list = [file for file in glob.glob(os.path.join(FOLDER_PATH, "*")) if
os.path.splitext(file)[1] not in IGNORED_EXTENSIONS and not os.path.isdir(file)]
print(f"Found {len(file_list)} files")
def convert_dicom_video(in_file_path, dicom, blank_proportion=None):
fourcc = cv2.VideoWriter_fourcc(*'MP4V')
out_file_path = f"{in_file_path}.mp4"
fps = 1000 / dicom.FrameTime
if len(dicom.pixel_array.shape) == 4:
n_frames, width, height, channels = dicom.pixel_array.shape
print(
f"{os.path.basename(in_file_path)} -> {os.path.basename(out_file_path)} / FPS: {fps} / WIDTH: {width} / HEIGHT: {height} / CHANNELS: {channels}")
else:
(n_frames, width, height), channels = dicom.pixel_array.shape, 1
print(
f"{os.path.basename(in_file_path)} -> {os.path.basename(out_file_path)} / FPS: {fps} / WIDTH: {width} / HEIGHT: {height} / CHANNELS: {channels}")
out = cv2.VideoWriter(out_file_path, fourcc, fps, (height, width), True)
for i, frame in enumerate(dicom.pixel_array):
if blank_proportion:
row_to = int(height * blank_proportion)
frame[0:row_to, :] = 0
if channels == 1:
frame = cv2.cvtColor(frame, cv2.COLOR_GRAY2RGB)
else:
if CONVERT_TO_BGR:
frame = frame[..., ::-1]
out.write(frame)
out.release()
cv2.destroyAllWindows()
def convert_dicom_frame(in_file_path, dicom, blank_proportion=None):
out_file_path = f"{in_file_path}.jpg"
if len(dicom.pixel_array.shape) == 2:
width, height = dicom.pixel_array.shape
frame = dicom.pixel_array
else:
width, height, channels = dicom.pixel_array.shape
frame = dicom.pixel_array
if CONVERT_TO_BGR:
frame = frame[..., ::-1]
if blank_proportion:
row_to = int(height * blank_proportion)
frame[0:row_to, :] = 0
print(f"{os.path.basename(in_file_path)} -> {os.path.basename(out_file_path)} / WIDTH: {width} / HEIGHT: {height}")
cv2.imwrite(out_file_path, frame)
if __name__ == "__main__":
failed_files = []
for file_path in file_list:
try:
dicom = pydicom.dcmread(file_path)
except pydicom.errors.InvalidDicomError:
failed_files.append(file_path)
continue
if type(dicom) == DicomDir:
failed_files.append(file_path)
else:
try:
if len(dicom.pixel_array.shape) == 4 or (
len(dicom.pixel_array.shape) == 3 and dicom.pixel_array.shape[2] != 3):
convert_dicom_video(file_path, dicom, blank_proportion=BLANK_PROPORTION)
elif len(dicom.pixel_array.shape) == 2 or (
len(dicom.pixel_array.shape) == 3 and dicom.pixel_array.shape[2] == 3):
convert_dicom_frame(file_path, dicom, blank_proportion=BLANK_PROPORTION)
else:
raise ValueError(f"Got shape: {dicom.pixel_array.shape}")
except AttributeError:
failed_files.append(file_path)
newline = '\n'
print(
f"Completed converting {len(file_list)-len(failed_files)} of {len(file_list)} files; the following {len(failed_files)} files failed:\n{newline.join(failed_files)}")
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment