Merge pull request #197 from databricks-industry-solutions/feature/remove_un_tags

Enhance DICOM metadata extraction by removing UN VR elements

Author: erinaldidb

README.md

@@ -165,6 +165,27 @@ metadata_df = DicomMetaAnonymizerExtractor(
 By setting the `anonym_mode` parameter to `"METADATA"`, the DICOM metadata will be anonymized during the ingestion process. This ensures that sensitive patient information is not stored in the catalog.
 The default configuration will save the anonymized DICOM files under `anonymization_base_path` property's path.
 
+## Remove UN Tags
+DICOM files can contain elements with Value Representation `UN` (Unknown), which are tags that could not be resolved to a specific VR during parsing. These tags often carry unstructured or proprietary data that can bloat the extracted metadata, cause serialization issues, or introduce noise in downstream analytics.
+
+Pixels provides a built-in option to strip all `UN` VR elements from the dataset before metadata extraction. The removal is recursive, so `UN` elements nested inside sequences (`SQ`) are also cleaned up.
+
+To enable this feature, set `remove_un_tags=True` on the `DicomMetaExtractor`:
+```python
+from dbx.pixels import Catalog
+from dbx.pixels.dicom import *
+
+catalog = Catalog(spark)
+catalog_df = catalog.catalog(<path>)
+
+meta_df = DicomMetaExtractor(
+    catalog,
+    remove_un_tags=True
+).transform(catalog_df)
+
+catalog.save(meta_df)
+```
+
 ---
 ## OHIF Viewer
 Inside `dbx.pixels` resources folder, a pre-built version of [OHIF Viewer](https://github.com/OHIF/Viewers) with Databricks and [Unity Catalog Volumes](https://docs.databricks.com/en/sql/language-manual/sql-ref-volumes.html) extension is provided. 

dbx/pixels/dicom/dicom_meta_extractor.py

@@ -18,7 +18,7 @@ class DicomMetaExtractor(Transformer):
     headers over the network, maximizing I/O throughput on each Spark task.
     """
 
-    MAX_WORKERS = 20
+    MAX_WORKERS = 32
 
     def __init__(
         self,
@@ -29,6 +29,7 @@ def __init__(
         deep=False,
         useVariant=True,
         maxWorkers=None,
+        remove_un_tags=False,
     ):
         self.inputCol = inputCol
         self.outputCol = outputCol
@@ -37,6 +38,7 @@ def __init__(
         self.deep = deep
         self.useVariant = useVariant
         self.maxWorkers = maxWorkers if maxWorkers is not None else self.MAX_WORKERS
+        self.remove_un_tags = remove_un_tags
 
     def check_input_type(self, schema):
         field = schema[self.inputCol]
@@ -69,6 +71,7 @@ def _transform(self, df):
         output_col = self.outputCol
         deep = self.deep
         max_workers = self.maxWorkers
+        remove_un_tags = self.remove_un_tags
 
         # Build output schema: all existing columns + the new meta column (as StringType initially)
         out_schema = t.StructType(
@@ -80,12 +83,12 @@ def _extract_meta(iterator: Iterator[pd.DataFrame]) -> Iterator[pd.DataFrame]:
             import simplejson as json
             from pydicom import dcmread
 
-            def _process_file(path: str, is_deep: bool, anon: bool) -> str:
+            def _process_file(path: str, deep: bool, anon: bool, remove_un_tags: bool) -> str:
                 try:
                     fp, fsize = cloud_open(path, anon)
-                    with dcmread(fp, defer_size=1000, stop_before_pixels=(not is_deep)) as dataset:
-                        meta_js = extract_metadata(dataset, is_deep)
-                        if is_deep:
+                    with dcmread(fp, defer_size=1000, stop_before_pixels=(not deep)) as dataset:
+                        meta_js = extract_metadata(dataset, deep, remove_un_tags=remove_un_tags)
+                        if deep:
                             meta_js["hash"] = hashlib.sha1(fp.read()).hexdigest()
                         meta_js["file_size"] = fsize
                         return json.dumps(meta_js, ignore_nan=True)
@@ -107,7 +110,7 @@ def _process_file(path: str, is_deep: bool, anon: bool) -> str:
                 with ThreadPoolExecutor(max_workers=max_workers) as executor:
                     meta_results = list(
                         executor.map(
-                            lambda args: _process_file(args[0], deep, args[1]),
+                            lambda args: _process_file(args[0], deep, args[1], remove_un_tags),
                             zip(paths, anon_flags),
                         )
                     )

dbx/pixels/dicom/dicom_utils.py

@@ -46,12 +46,38 @@ def check_pixel_data(ds: Dataset) -> Dataset | None:
         return None
 
 
-def extract_metadata(ds: Dataset, deep: bool = True) -> dict:
+def remove_un_vr_elements(ds: Dataset) -> Dataset:
+    """
+    Recursively removes all elements with VR='UN' from a pydicom dataset.
+    """
+    # Create a list of tags to delete so we don't modify the dataset while iterating
+    tags_to_delete = []
+
+    for elem in ds:
+        if elem.VR == "UN":
+            tags_to_delete.append(elem.tag)
+        elif elem.VR == "SQ":
+            # If the element is a sequence, recursively process each dataset within it
+            for item in elem:
+                remove_un_vr_elements(item)
+
+    # Delete the identified tags from the current dataset level
+    for tag in tags_to_delete:
+        del ds[tag]
+
+    return ds
+
+
+def extract_metadata(ds: Dataset, deep: bool = True, remove_un_tags: bool = False) -> dict:
     """Extract metadata from header of dicom image file
     params:
       path -- local path like /dbfs/mnt/... or s3://<bucket>/path/to/object.dcm
       deep -- True if deep inspection of the Dicom header is required
+      remove_un_vr_elements -- True if UN VR elements should be removed from the dataset
     """
+    if remove_un_tags:
+        ds = remove_un_vr_elements(ds)
+
     a = None
     js = {}