scoliosis-inference/io_dicom.py

import os
from typing import Optional

import cv2
import numpy as np
import pydicom

def read_dicom(path: str) -> np.ndarray:
    ds = pydicom.dcmread(path)
    img = ds.pixel_array

    if getattr(ds, "PhotometricInterpretation", "") == "MONOCHROME1":
        img = np.max(img) - img

    img = img.astype(np.float32)
    img -= img.min()
    if img.max() > 0:
        img /= img.max()
    img = (img * 255.0).clip(0, 255).astype(np.uint8)

    if img.ndim == 2:
        img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)
    elif img.shape[2] == 3:
        img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)

    return img


def save_secondary_capture(
    original_dicom_path: str,
    img_bgr: np.ndarray,
    out_path: str,
    series_description: str = "RG-SCOLIOSIS-CONTOUR",
    overlay_text: Optional[str] = None,
):
    import pydicom
    from pydicom.uid import generate_uid

    ds = pydicom.dcmread(original_dicom_path)

    img_draw = img_bgr.copy()
    if overlay_text:
        cv2.putText(
            img_draw,
            overlay_text,
            (20, 40),
            cv2.FONT_HERSHEY_SIMPLEX,
            1.0,
            (0, 255, 0),
            2,
        )

    img_rgb = cv2.cvtColor(img_draw, cv2.COLOR_BGR2RGB)
    rows, cols = img_rgb.shape[:2]

    ds.SOPClassUID = "1.2.840.10008.5.1.4.1.1.7" 
    ds.SOPInstanceUID = generate_uid()
    ds.SeriesInstanceUID = generate_uid()
    ds.Modality = "OT"
    ds.SeriesDescription = series_description

    if "SeriesNumber" in ds:
        try:
            ds.SeriesNumber = int(ds.SeriesNumber) + 1000
        except Exception:
            ds.SeriesNumber = 1000
    else:
        ds.SeriesNumber = 1000

    ds.Rows = rows
    ds.Columns = cols
    ds.SamplesPerPixel = 3
    ds.PhotometricInterpretation = "RGB"
    ds.PlanarConfiguration = 0
    ds.BitsAllocated = 8
    ds.BitsStored = 8
    ds.HighBit = 7
    ds.PixelRepresentation = 0
    ds.PixelData = img_rgb.tobytes()

    os.makedirs(os.path.dirname(out_path), exist_ok=True)
    ds.save_as(out_path, write_like_original=False)


def save_dicom_sr(
    original_dicom_path: str,
    description_text: str,
    conclusion_text: str,
    out_path: str,
    series_description: str = "SCOLIOSIS_SR",
    extra_struct: Optional[dict] = None,
):
    import datetime as _dt
    import pydicom
    from pydicom.dataset import Dataset, FileDataset
    from pydicom.uid import ExplicitVRLittleEndian, generate_uid

    src = pydicom.dcmread(original_dicom_path)

    file_meta = Dataset()
    file_meta.MediaStorageSOPClassUID = "1.2.840.10008.5.1.4.1.1.88.11"
    file_meta.MediaStorageSOPInstanceUID = generate_uid()
    file_meta.TransferSyntaxUID = ExplicitVRLittleEndian
    file_meta.ImplementationClassUID = generate_uid()

    ds = FileDataset(out_path, {}, file_meta=file_meta, preamble=b"\0" * 128)

    for tag in (
        "PatientName",
        "PatientID",
        "PatientBirthDate",
        "PatientSex",
        "StudyInstanceUID",
        "StudyID",
        "AccessionNumber",
        "StudyDate",
        "StudyTime",
    ):
        if tag in src:
            ds[tag] = src[tag]

    ds.SOPClassUID = file_meta.MediaStorageSOPClassUID
    ds.SOPInstanceUID = file_meta.MediaStorageSOPInstanceUID
    ds.SeriesInstanceUID = generate_uid()
    ds.Modality = "SR"
    ds.SeriesDescription = series_description
    ds.SeriesNumber = 2000
    ds.InstanceNumber = 1

    if "SpecificCharacterSet" in src:
        ds.SpecificCharacterSet = src.SpecificCharacterSet
    else:
        ds.SpecificCharacterSet = "ISO_IR 192"

    now = _dt.datetime.now()
    ds.ContentDate = now.strftime("%Y%m%d")
    ds.ContentTime = now.strftime("%H%M%S")

    root = Dataset()
    root.ValueType = "CONTAINER"
    root.ContinuityOfContent = "SEPARATE"
    root.ConceptNameCodeSequence = [Dataset()]
    root.ConceptNameCodeSequence[0].CodeValue = "11528-7"
    root.ConceptNameCodeSequence[0].CodingSchemeDesignator = "LN"
    root.ConceptNameCodeSequence[0].CodeMeaning = "Imaging Report"

    desc_item = Dataset()
    desc_item.ValueType = "TEXT"
    desc_item.ConceptNameCodeSequence = [Dataset()]
    desc_item.ConceptNameCodeSequence[0].CodeValue = "121070"
    desc_item.ConceptNameCodeSequence[0].CodingSchemeDesignator = "DCM"
    desc_item.ConceptNameCodeSequence[0].CodeMeaning = "Findings"
    desc_item.TextValue = str(description_text)

    concl_item = Dataset()
    concl_item.ValueType = "TEXT"
    concl_item.ConceptNameCodeSequence = [Dataset()]
    concl_item.ConceptNameCodeSequence[0].CodeValue = "121106"
    concl_item.ConceptNameCodeSequence[0].CodingSchemeDesignator = "DCM"
    concl_item.ConceptNameCodeSequence[0].CodeMeaning = "Summary"
    concl_item.TextValue = str(conclusion_text)

    content_items = [desc_item, concl_item]

    if extra_struct:
        def _text_item(code, meaning, value, scheme="DCM"):
            item = Dataset()
            item.ValueType = "TEXT"
            item.ConceptNameCodeSequence = [Dataset()]
            item.ConceptNameCodeSequence[0].CodeValue = code
            item.ConceptNameCodeSequence[0].CodingSchemeDesignator = scheme
            item.ConceptNameCodeSequence[0].CodeMeaning = meaning
            item.TextValue = str(value)
            return item

        if "degree_class" in extra_struct:
            content_items.append(_text_item("SCL-DEGREE", "Scoliosis degree class (1-4, 0=none)", extra_struct["degree_class"]))
        if "scoliosis_type" in extra_struct:
            content_items.append(_text_item("SCL-TYPE", "Scoliosis type", extra_struct["scoliosis_type"]))
        if "main_cobb" in extra_struct:
            content_items.append(_text_item("SCL-MAIN-COBB", "Main Cobb angle, deg", f"{extra_struct['main_cobb']:.1f}"))
        if "prob_scoliosis" in extra_struct:
            content_items.append(_text_item("SCL-PROB", "Probability of scoliosis (0-1)", f"{extra_struct['prob_scoliosis']:.2f}"))

    root.ContentSequence = content_items
    ds.ContentSequence = [root]
    ds.CompletionFlag = "COMPLETE"
    ds.VerificationFlag = "UNVERIFIED"

    os.makedirs(os.path.dirname(out_path), exist_ok=True)
    ds.save_as(out_path, write_like_original=False)