Visual_NLP_Metrics.ipynb

{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "98MPFXW9vXro"
      },
      "outputs": [],
      "source": [
        "!python /content/prepare_data.py"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 306
        },
        "id": "4Bczx-9qu9kD",
        "outputId": "e35b968f-58a4-4b6d-bcf5-b0651d58bda8"
      },
      "outputs": [],
      "source": [
        "from sparkocr import start\n",
        "import os\n",
        "import json\n",
        "import time\n",
        "import shutil\n",
        "\n",
        "import json\n",
        "\n",
        "with open(\"./creds.json\", \"r\") as creds_in:\n",
        "    creds = json.loads(creds_in.read())[\"Credentials\"]\n",
        "\n",
        "os.environ['AWS_ACCESS_KEY_ID'] = creds[\"AccessKeyId\"]\n",
        "os.environ['AWS_SECRET_ACCESS_KEY'] = creds[\"SecretAccessKey\"]\n",
        "os.environ['AWS_SESSION_TOKEN'] = creds[\"SessionToken\"]\n",
        "os.environ['SPARK_OCR_LICENSE'] = creds[\"SPARK_OCR_LICENSE\"]\n",
        "\n",
        "colab_max = {\n",
        "    \"spark.driver.memory\": \"10g\",\n",
        "    \"spark.executor.memory\": \"12g\",\n",
        "    \"spark.executor.memoryOverhead\": \"3g\",\n",
        "    \"spark.driver.memoryOverhead\": \"2g\",\n",
        "    \"spark.extraListeners\": \"com.johnsnowlabs.license.LicenseLifeCycleManager\",\n",
        "    \"spark.cores.max\": \"12\",\n",
        "    \"spark.executor.cores\": \"4\"\n",
        "}\n",
        "\n",
        "spark = start(jar_path=\"./\",\n",
        "              nlp_secret=\"5.5.2.PR-2579.6fce19f4d0e8cce1d61a808968c833c932bf2dd1\",\n",
        "              nlp_internal=True,\n",
        "              nlp_jsl=True,\n",
        "              nlp_version=\"5.5.2\",\n",
        "              extra_conf=colab_max,\n",
        "              use_gpu=True)\n",
        "\n",
        "spark"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "oYdEFqUcvBcU",
        "outputId": "f184d751-0dbb-4396-ef94-316d93ce6935"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "import json\n",
        "import sys\n",
        "import time\n",
        "from PIL import Image\n",
        "import io\n",
        "\n",
        "from sparknlp.annotator import *\n",
        "from sparknlp.base import *\n",
        "import sparknlp_jsl\n",
        "from sparknlp_jsl.annotator import *\n",
        "\n",
        "import sparkocr\n",
        "from sparkocr.transformers import *\n",
        "from sparkocr.utils import *\n",
        "from sparkocr.enums import *\n",
        "\n",
        "from pyspark.ml import PipelineModel, Pipeline\n",
        "from sparknlp.pretrained import PretrainedPipeline\n",
        "from pyspark.ml import PipelineModel, Pipeline\n",
        "\n",
        "import pyspark.sql.functions as F"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ghnCMYNzvPwp"
      },
      "outputs": [],
      "source": [
        "def calculate_precision(gt, pred):\n",
        "    tp = [i for i in pred if i in gt]\n",
        "    try:\n",
        "        precision = len(tp) / len(pred)\n",
        "    except ZeroDivisionError:\n",
        "        precision = 0\n",
        "\n",
        "    return precision\n",
        "\n",
        "def calculate_recall(gt, pred):\n",
        "    tp = [i for i in pred if i in gt]\n",
        "    try:\n",
        "        recall = len(tp) / len(gt)\n",
        "    except ZeroDivisionError:\n",
        "        recall = 0\n",
        "\n",
        "    return recall\n",
        "\n",
        "def calculate_metrics(gt_path, df_path, save_result_path, ner_chunk=\"merged_ner_chunk\"):\n",
        "\n",
        "    \"\"\"\n",
        "    gt_path - The location of json file containing ground truth. Generate using prepare_data.py\n",
        "    df_path - The location of parquet file containing result from Visual NLP pipeline.\n",
        "    save_result_path - The location of the output json file containing result\n",
        "    ner_chunk - The column name containing final chunks used for metrics.\n",
        "    \"\"\"\n",
        "\n",
        "    total_precision = 0.0\n",
        "    total_recall = 0.0\n",
        "\n",
        "    with open(gt_path, \"r\") as file_in:\n",
        "        gt = json.loads(file_in.read())\n",
        "\n",
        "    df = spark.read.format(\"parquet\").load(df_path).select(\"path\", ner_chunk)\n",
        "\n",
        "    total_files = df.count()\n",
        "    print(f\"Total Files : {total_files}\")\n",
        "\n",
        "    save_result = []\n",
        "\n",
        "    for item_row in df.toLocalIterator():\n",
        "        data = item_row.asDict()\n",
        "\n",
        "        img_path = os.path.basename(data[\"path\"])\n",
        "        chunks = data[ner_chunk]\n",
        "\n",
        "        collect_pred = []\n",
        "        for chunk in chunks:\n",
        "            pred = chunk.asDict()[\"result\"].strip().lower()\n",
        "            collect_pred.append(pred)\n",
        "\n",
        "        lowered_gt = [i.lower() for i in gt[img_path]]\n",
        "\n",
        "        precision = calculate_precision(lowered_gt, collect_pred)\n",
        "        recall = calculate_recall(lowered_gt, collect_pred)\n",
        "\n",
        "        total_precision += precision\n",
        "        total_recall += recall\n",
        "\n",
        "        save_result.append({\"Filename\" : img_path,\n",
        "                             \"Ground Truth\" : lowered_gt,\n",
        "                             \"Predicted PHI\" : collect_pred,\n",
        "                             \"Precision\" : precision,\n",
        "                             \"Recall\" : recall})\n",
        "\n",
        "    total_precision = round(total_precision / total_files, 3)\n",
        "    total_recall = round(total_recall / total_files, 3)\n",
        "    f1_score = round(2 * ((total_precision * total_recall) / (total_precision + total_recall)), 3)\n",
        "\n",
        "    print(f\"Precision : {total_precision}\")\n",
        "    print(f\"Recall : {total_recall}\")\n",
        "    print(f\"F1-Score : {f1_score}\")\n",
        "\n",
        "    with open(save_result_path, \"w\") as result_file_out:\n",
        "        json.dump(save_result, result_file_out, indent=4)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "6qYjS66wvI6k"
      },
      "outputs": [],
      "source": [
        "os.makedirs(\"./results/\", exist_ok=True)\n",
        "os.makedirs(\"./results/deid_image_result/\", exist_ok=True)\n",
        "os.makedirs(\"./results/deid_image_result/visual_nlp/\", exist_ok=True)\n",
        "os.makedirs(\"./results/deid_image_result/presidio/\", exist_ok=True)\n",
        "os.makedirs(\"./results/ner_result/\", exist_ok=True)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "QP8UZIfMvLLg",
        "outputId": "e4769426-b818-4e49-ca10-c9b5f88881ef"
      },
      "outputs": [],
      "source": [
        "dicom_to_image = DicomToImageV3() \\\n",
        "    .setInputCols([\"content\"]) \\\n",
        "    .setOutputCol(\"image_raw\") \\\n",
        "    .setKeepInput(False)\n",
        "\n",
        "documentAssembler = DocumentAssembler() \\\n",
        "    .setInputCol(\"text\") \\\n",
        "    .setOutputCol(\"document_raw\") \\\n",
        "    .setCleanupMode(\"shrink\")\n",
        "\n",
        "cleanUpPatterns = [\"<[^>]*>\", r\"\\^ ?\", r\"\\^\"]\n",
        "documentNormalizer = DocumentNormalizer() \\\n",
        "  .setInputCols(\"document_raw\") \\\n",
        "  .setOutputCol(\"document\") \\\n",
        "  .setAction(\"clean\") \\\n",
        "  .setPatterns(cleanUpPatterns) \\\n",
        "  .setReplacement(\" \") \\\n",
        "  .setPolicy(\"pretty_all\")\n",
        "\n",
        "sentencerDL = SentenceDetectorDLModel\\\n",
        "    .pretrained(\"sentence_detector_dl\", \"en\") \\\n",
        "    .setInputCols([\"document\"]) \\\n",
        "    .setOutputCol(\"sentence\")\n",
        "\n",
        "tokenizer = Tokenizer() \\\n",
        "    .setInputCols([\"sentence\"]) \\\n",
        "    .setOutputCol(\"token\")\n",
        "\n",
        "word_embeddings = WordEmbeddingsModel.pretrained(\"embeddings_clinical\", \"en\", \"clinical/models\") \\\n",
        "    .setInputCols([\"sentence\", \"token\"]) \\\n",
        "    .setOutputCol(\"embeddings\")\n",
        "\n",
        "clinical_ner = MedicalNerModel.pretrained(\"ner_deid_large\", \"en\", \"clinical/models\") \\\n",
        "    .setInputCols([\"sentence\", \"token\", \"embeddings\"]) \\\n",
        "    .setOutputCol(\"ner_subentity\")\n",
        "\n",
        "regex_matcher = RegexMatcher()\\\n",
        "    .setInputCols(\"sentence\")\\\n",
        "    .setOutputCol(\"regex\")\\\n",
        "    .setRules([\"[0-9]{4} (JAN|FEB|MAR|APR|MAY|JUN|JUL|AUG|SEP|OCT|NOV|DEC) [0-9]{1,2};DATE\",\n",
        "               \"(0[1-9]|1[0-2])/(0[1-9]|[12][0-9]|3[01])/\\d{4};DATENUM\",\n",
        "               \"(?:\\s[MF]|\\b[MF])(?:\\s|\\b|$);GENDER\",\n",
        "               \"(?<!\\d)\\d{9,10}(?!\\d);ID\"])\\\n",
        "    .setDelimiter(\";\")\n",
        "\n",
        "chunkConverter = ChunkConverter()\\\n",
        "    .setInputCols(\"regex\")\\\n",
        "    .setOutputCol(\"regex_chunks\")\n",
        "\n",
        "custom_ner_converter_internal = NerConverterInternalModel() \\\n",
        "    .setInputCols([\"sentence\",\"token\", \"ner_subentity\"]) \\\n",
        "    .setOutputCol(\"ner_chunk\") \\\n",
        "    .setThreshold(0.4) \\\n",
        "    .setWhiteList(['NAME', 'AGE', 'LOCATION', 'PERSON', 'DOCTOR', 'PATIENT'])\n",
        "\n",
        "chunk_merger = ChunkMergeApproach()\\\n",
        "    .setInputCols('regex_chunks', \"ner_chunk\")\\\n",
        "    .setOutputCol('merged_ner_chunk')\\\n",
        "    .setMergeOverlapping(True)\n",
        "\n",
        "base_stages = [\n",
        "    dicom_to_image,\n",
        "    documentAssembler,\n",
        "    documentNormalizer,\n",
        "    sentencerDL,\n",
        "    tokenizer,\n",
        "    word_embeddings,\n",
        "    clinical_ner,\n",
        "    regex_matcher,\n",
        "    chunkConverter,\n",
        "    custom_ner_converter_internal,\n",
        "    chunk_merger\n",
        "]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "CMVw_3L4vV3f",
        "outputId": "382a3d48-e488-4771-fdb6-e1692351a135"
      },
      "outputs": [],
      "source": [
        "text_detector = ImageTextDetector.pretrained(\"image_text_detector_mem_opt\", \"en\", \"clinical/ocr\") \\\n",
        "    .setInputCol(\"image_raw\") \\\n",
        "    .setOutputCol(\"text_regions\") \\\n",
        "    .setScoreThreshold(0.7) \\\n",
        "    .setLinkThreshold(0.5) \\\n",
        "    .setWithRefiner(True) \\\n",
        "    .setTextThreshold(0.4) \\\n",
        "    .setSizeThreshold(-1) \\\n",
        "    .setUseGPU(True) \\\n",
        "    .setWidth(0)\n",
        "\n",
        "ocr = ImageToTextV2.pretrained(\"ocr_base_printed_v2_opt\", \"en\", \"clinical/ocr\") \\\n",
        "    .setRegionsColumn(\"text_regions\") \\\n",
        "    .setInputCols([\"image_raw\"]) \\\n",
        "    .setOutputCol(\"text\") \\\n",
        "    .setOutputFormat(\"text_with_positions\") \\\n",
        "    .setGroupImages(False) \\\n",
        "    .setKeepInput(False) \\\n",
        "    .setUseGPU(True) \\\n",
        "    .setUseCaching(True) \\\n",
        "    .setBatchSize(4)\n",
        "\n",
        "new_stages = base_stages.copy()\n",
        "\n",
        "new_stages.insert(1,text_detector)\n",
        "new_stages.insert(2,ocr)\n",
        "\n",
        "detector_v2_base_v2 = Pipeline(stages=new_stages)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "px-9fQbTv2yO",
        "outputId": "00ba8d36-6953-44be-dacf-d77a1de8c4e0"
      },
      "outputs": [],
      "source": [
        "df = spark.read.format(\"binaryFile\").load(\"./dicom_files/*.dcm\")\n",
        "\n",
        "result_base = detector_v2_base_v2.fit(df).transform(df)\n",
        "\n",
        "start = time.time()\n",
        "\n",
        "result_base.write.format(\"parquet\").mode(\"overwrite\").save(\"./result_base/\")\n",
        "\n",
        "end = time.time()\n",
        "\n",
        "print(f\"Time Taken : {round((end - start) / df.count(), 2)}\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "6BMdDO6Xv6yn",
        "outputId": "893c0f7e-7d28-4ef4-ec45-47429da21362"
      },
      "outputs": [],
      "source": [
        "gt_path = \"./gt_ocr.json\"\n",
        "df_path = \"./result_base/\"\n",
        "result_path = \"./results/ner_result/result_v2_base.json\"\n",
        "\n",
        "calculate_metrics(gt_path, df_path, result_path)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "3lpV4rY8zh46",
        "outputId": "545de60a-c29d-4c8a-b2c8-3afeb644e1c7"
      },
      "outputs": [],
      "source": [
        "text_detector = ImageTextDetector.pretrained(\"image_text_detector_mem_opt\", \"en\", \"clinical/ocr\") \\\n",
        "    .setInputCol(\"image_raw\") \\\n",
        "    .setOutputCol(\"text_regions\") \\\n",
        "    .setScoreThreshold(0.7) \\\n",
        "    .setLinkThreshold(0.5) \\\n",
        "    .setWithRefiner(True) \\\n",
        "    .setTextThreshold(0.4) \\\n",
        "    .setSizeThreshold(-1) \\\n",
        "    .setUseGPU(True) \\\n",
        "    .setWidth(0)\n",
        "\n",
        "ocr = ImageToTextV2.pretrained(\"ocr_large_printed_v2_opt\", \"en\", \"clinical/ocr\") \\\n",
        "    .setRegionsColumn(\"text_regions\") \\\n",
        "    .setInputCols([\"image_raw\"]) \\\n",
        "    .setOutputCol(\"text\") \\\n",
        "    .setOutputFormat(\"text_with_positions\") \\\n",
        "    .setGroupImages(False) \\\n",
        "    .setKeepInput(False) \\\n",
        "    .setUseGPU(True) \\\n",
        "    .setUseCaching(True) \\\n",
        "    .setBatchSize(4)\n",
        "\n",
        "new_stages = base_stages.copy()\n",
        "\n",
        "new_stages.insert(1,text_detector)\n",
        "new_stages.insert(2,ocr)\n",
        "\n",
        "detector_v2_large_v2 = Pipeline(stages=new_stages)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "pPny_XYWzlLK",
        "outputId": "f534bd15-0436-4e22-9667-a8a4c2fb2c86"
      },
      "outputs": [],
      "source": [
        "df = spark.read.format(\"binaryFile\").load(\"./dicom_files/*.dcm\")\n",
        "\n",
        "result_large = detector_v2_large_v2.fit(df).transform(df)\n",
        "\n",
        "start = time.time()\n",
        "\n",
        "result_large.write.format(\"parquet\").mode(\"overwrite\").save(\"./result_large/\")\n",
        "\n",
        "end = time.time()\n",
        "\n",
        "print(f\"Time Taken : {round((end - start) / df.count(), 2)}\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "dfvUkrQt5OGW",
        "outputId": "5ea41f69-fc2e-486c-dfc6-357e8ffc1a38"
      },
      "outputs": [],
      "source": [
        "gt_path = \"./gt_ocr.json\"\n",
        "df_path = \"./result_large/\"\n",
        "result_path = \"./results/ner_result/result_v2_large.json\"\n",
        "\n",
        "calculate_metrics(gt_path, df_path, result_path)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "BFa_pK8x5gd5",
        "outputId": "70a5dbf3-cd78-4da5-a058-1cd372b15648"
      },
      "outputs": [],
      "source": [
        "text_detector = ImageTextDetector.pretrained(\"image_text_detector_mem_opt\", \"en\", \"clinical/ocr\") \\\n",
        "    .setInputCol(\"image_raw\") \\\n",
        "    .setOutputCol(\"text_regions\") \\\n",
        "    .setScoreThreshold(0.7) \\\n",
        "    .setLinkThreshold(0.5) \\\n",
        "    .setWithRefiner(True) \\\n",
        "    .setTextThreshold(0.4) \\\n",
        "    .setSizeThreshold(-1) \\\n",
        "    .setUseGPU(True) \\\n",
        "    .setWidth(0)\n",
        "\n",
        "ocr = ImageToTextV3() \\\n",
        "    .setInputCols([\"image_raw\", \"text_regions\"]) \\\n",
        "    .setOutputCol(\"text\")\n",
        "\n",
        "new_stages = base_stages.copy()\n",
        "\n",
        "new_stages.insert(1,text_detector)\n",
        "new_stages.insert(2,ocr)\n",
        "\n",
        "detector_v3 = Pipeline(stages=new_stages)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "YlfUcyk76DL1",
        "outputId": "bcfd9228-5e61-45cd-90c9-a16d307af1b8"
      },
      "outputs": [],
      "source": [
        "df = spark.read.format(\"binaryFile\").load(\"./dicom_files/*.dcm\")\n",
        "\n",
        "result_v3 = detector_v3.fit(df).transform(df)\n",
        "\n",
        "start = time.time()\n",
        "\n",
        "result_v3.write.format(\"parquet\").mode(\"overwrite\").save(\"./result_v3/\")\n",
        "\n",
        "end = time.time()\n",
        "\n",
        "print(f\"Time Taken : {round((end - start) / df.count(), 2)}\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "pHfzqJPR6Gk1",
        "outputId": "691888a5-9d9e-430d-c6d6-ccb5d6a4900b"
      },
      "outputs": [],
      "source": [
        "gt_path = \"./gt_ocr.json\"\n",
        "df_path = \"./result_v3/\"\n",
        "result_path = \"./results/ner_result/result_v3.json\"\n",
        "\n",
        "calculate_metrics(gt_path, df_path, result_path)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "oQTYaJGT6HNi"
      },
      "outputs": [],
      "source": [
        "ocr = ImageToText() \\\n",
        "    .setInputCol(\"image_raw\") \\\n",
        "    .setOutputCol(\"text\") \\\n",
        "    .setIgnoreResolution(False) \\\n",
        "    .setPageIteratorLevel(PageIteratorLevel.SYMBOL) \\\n",
        "    .setPageSegMode(PageSegmentationMode.SPARSE_TEXT) \\\n",
        "    .setWithSpaces(True) \\\n",
        "    .setConfidenceThreshold(70)\n",
        "\n",
        "new_stages = base_stages.copy()\n",
        "\n",
        "new_stages.insert(1,ocr)\n",
        "\n",
        "tesseract_pipe = Pipeline(stages=new_stages)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "C-vJOV2m9DxG",
        "outputId": "cbce6241-298d-4337-a6c3-e4b3ea9342a7"
      },
      "outputs": [],
      "source": [
        "df = spark.read.format(\"binaryFile\").load(\"./dicom_files/*.dcm\")\n",
        "\n",
        "result_tesseract = tesseract_pipe.fit(df).transform(df)\n",
        "\n",
        "start = time.time()\n",
        "\n",
        "result_tesseract.write.format(\"parquet\").mode(\"overwrite\").save(\"./result_tesseract/\")\n",
        "\n",
        "end = time.time()\n",
        "\n",
        "print(f\"Time Taken : {round((end - start) / df.count(), 2)}\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "19Yy6I6T9Fdl",
        "outputId": "9224ca03-4277-46fd-f1ec-1966525590d8"
      },
      "outputs": [],
      "source": [
        "gt_path = \"./gt_ocr.json\"\n",
        "df_path = \"./result_tesseract/\"\n",
        "result_path = \"./results/ner_result/result_tesseract.json\"\n",
        "\n",
        "calculate_metrics(gt_path, df_path, result_path)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "uNMpB7PD9Hdo",
        "outputId": "8f723f21-babc-43a4-bfa5-bffdda16f243"
      },
      "outputs": [],
      "source": [
        "dicom_to_image = DicomToImageV3() \\\n",
        "    .setInputCols([\"content\"]) \\\n",
        "    .setOutputCol(\"image_raw\") \\\n",
        "    .setKeepInput(False)\n",
        "\n",
        "documentAssembler = DocumentAssembler() \\\n",
        "    .setInputCol(\"text\") \\\n",
        "    .setOutputCol(\"document_raw\") \\\n",
        "    .setCleanupMode(\"shrink\")\n",
        "\n",
        "cleanUpPatterns = [\"<[^>]*>\", r\"\\^ ?\", r\"\\^\"]\n",
        "documentNormalizer = DocumentNormalizer() \\\n",
        "  .setInputCols(\"document_raw\") \\\n",
        "  .setOutputCol(\"document\") \\\n",
        "  .setAction(\"clean\") \\\n",
        "  .setPatterns(cleanUpPatterns) \\\n",
        "  .setReplacement(\" \") \\\n",
        "  .setPolicy(\"pretty_all\")\n",
        "\n",
        "sentencerDL = SentenceDetectorDLModel\\\n",
        "    .pretrained(\"sentence_detector_dl\", \"en\") \\\n",
        "    .setInputCols([\"document\"]) \\\n",
        "    .setOutputCol(\"sentence\")\n",
        "\n",
        "tokenizer = Tokenizer() \\\n",
        "    .setInputCols([\"sentence\"]) \\\n",
        "    .setOutputCol(\"token\")\n",
        "\n",
        "word_embeddings = WordEmbeddingsModel.pretrained(\"embeddings_clinical\", \"en\", \"clinical/models\") \\\n",
        "    .setInputCols([\"sentence\", \"token\"]) \\\n",
        "    .setOutputCol(\"embeddings\")\n",
        "\n",
        "clinical_ner = MedicalNerModel.pretrained(\"ner_deid_large\", \"en\", \"clinical/models\") \\\n",
        "    .setInputCols([\"sentence\", \"token\", \"embeddings\"]) \\\n",
        "    .setOutputCol(\"ner_subentity\")\n",
        "\n",
        "regex_matcher = RegexMatcher()\\\n",
        "    .setInputCols(\"sentence\")\\\n",
        "    .setOutputCol(\"regex\")\\\n",
        "    .setRules([\"[0-9]{4} (JAN|FEB|MAR|APR|MAY|JUN|JUL|AUG|SEP|OCT|NOV|DEC) [0-9]{1,2};DATE\",\n",
        "               \"(0[1-9]|1[0-2])/(0[1-9]|[12][0-9]|3[01])/\\d{4};DATENUM\",\n",
        "               \"(?:\\s[MF]|\\b[MF])(?:\\s|\\b|$);GENDER\",\n",
        "               \"(?<!\\d)\\d{9,10}(?!\\d);ID\"])\\\n",
        "    .setDelimiter(\";\")\n",
        "\n",
        "chunkConverter = ChunkConverter()\\\n",
        "    .setInputCols(\"regex\")\\\n",
        "    .setOutputCol(\"regex_chunks\")\n",
        "\n",
        "custom_ner_converter_internal = NerConverterInternalModel() \\\n",
        "    .setInputCols([\"sentence\",\"token\", \"ner_subentity\"]) \\\n",
        "    .setOutputCol(\"ner_chunk\") \\\n",
        "    .setThreshold(0.4) \\\n",
        "    .setWhiteList(['NAME', 'AGE', 'LOCATION', 'PERSON', 'DOCTOR', 'PATIENT'])\n",
        "\n",
        "chunk_merger = ChunkMergeApproach()\\\n",
        "    .setInputCols('regex_chunks', \"ner_chunk\")\\\n",
        "    .setOutputCol('merged_ner_chunk')\\\n",
        "    .setMergeOverlapping(True)\n",
        "\n",
        "position_finder = PositionFinder() \\\n",
        "    .setInputCols(\"merged_ner_chunk\") \\\n",
        "    .setOutputCol(\"coordinates\") \\\n",
        "    .setPageMatrixCol(\"positions\") \\\n",
        "    .setOcrScaleFactor(0.9)\n",
        "\n",
        "draw_regions = DicomDrawRegions() \\\n",
        "    .setInputCol(\"path\") \\\n",
        "    .setInputRegionsCol(\"coordinates\") \\\n",
        "    .setOutputCol(\"dicom_cleaned\") \\\n",
        "    .setAggCols([\"path\"]) \\\n",
        "    .setKeepInput(True)\n",
        "\n",
        "draw_base_stages = [\n",
        "    dicom_to_image,\n",
        "    documentAssembler,\n",
        "    documentNormalizer,\n",
        "    sentencerDL,\n",
        "    tokenizer,\n",
        "    word_embeddings,\n",
        "    clinical_ner,\n",
        "    regex_matcher,\n",
        "    chunkConverter,\n",
        "    custom_ner_converter_internal,\n",
        "    chunk_merger,\n",
        "    position_finder,\n",
        "    draw_regions\n",
        "]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "2S14R5R_9dWY",
        "outputId": "26cbd0c2-4358-4b34-cf3e-0ec9bf8b941d"
      },
      "outputs": [],
      "source": [
        "text_detector = ImageTextDetector.pretrained(\"image_text_detector_mem_opt\", \"en\", \"clinical/ocr\") \\\n",
        "    .setInputCol(\"image_raw\") \\\n",
        "    .setOutputCol(\"text_regions\") \\\n",
        "    .setScoreThreshold(0.7) \\\n",
        "    .setLinkThreshold(0.5) \\\n",
        "    .setWithRefiner(True) \\\n",
        "    .setTextThreshold(0.4) \\\n",
        "    .setSizeThreshold(-1) \\\n",
        "    .setUseGPU(True) \\\n",
        "    .setWidth(0)\n",
        "\n",
        "ocr = ImageToTextV2.pretrained(\"ocr_large_printed_v2_opt\", \"en\", \"clinical/ocr\") \\\n",
        "    .setRegionsColumn(\"text_regions\") \\\n",
        "    .setInputCols([\"image_raw\"]) \\\n",
        "    .setOutputCol(\"text\") \\\n",
        "    .setOutputFormat(\"text_with_positions\") \\\n",
        "    .setGroupImages(False) \\\n",
        "    .setKeepInput(False) \\\n",
        "    .setUseGPU(True) \\\n",
        "    .setUseCaching(True) \\\n",
        "    .setBatchSize(4)\n",
        "\n",
        "new_stages = draw_base_stages.copy()\n",
        "\n",
        "new_stages.insert(1,text_detector)\n",
        "new_stages.insert(2,ocr)\n",
        "\n",
        "detector_v2_large_draw = Pipeline(stages=new_stages)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "background_save": true
        },
        "id": "7jvsEfek9iVx"
      },
      "outputs": [],
      "source": [
        "df = spark.read.format(\"binaryFile\").load(\"./dicom_files/*.dcm\")\n",
        "\n",
        "result_large_dicom = detector_v2_large_draw.fit(df).transform(df)\n",
        "\n",
        "result_large_dicom.write.format(\"parquet\").mode(\"overwrite\").save(\"./result_large_dicom/\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "OMiUeLOf9kSf"
      },
      "outputs": [],
      "source": [
        "import pydicom\n",
        "\n",
        "df = spark.read.format(\"parquet\").load(\"./result_large_dicom/\")\n",
        "\n",
        "for row in df.select(\"path\", \"dicom_cleaned\").toLocalIterator():\n",
        "    data = row.asDict()\n",
        "    dicom_file = data[\"dicom_cleaned\"]\n",
        "    path = data[\"path\"]\n",
        "\n",
        "    base_name = os.path.basename(path).replace(\".dcm\", \".jpg\")\n",
        "\n",
        "    ds = pydicom.dcmread(io.BytesIO(dicom_file))\n",
        "    image = Image.fromarray(ds.pixel_array).convert(\"L\")\n",
        "\n",
        "    file_out = os.path.join(\"./results/deid_image_result/visual_nlp/\", base_name)\n",
        "    image.save(file_out)"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "gpuType": "A100",
      "machine_shape": "hm",
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}