replace prepare script

mu_map/polar_map/prepare.py

@@ -2,7 +2,7 @@ import argparse
 import os
 from typing import List
 
-from mu_map.file.dicom import DICOM
+from mu_map.file.dicom import DICOM, DICOMTime
 
 
 def get_files_recursive(_dir: str) -> List[str]:
@@ -25,7 +25,7 @@ def get_files_recursive(_dir: str) -> List[str]:
 
 
 def is_scatter_corrected(dcm: DICOM) -> bool:
-    return not ("NoSC" in dcm.SeriesDescription)
+    return "SC " in dcm.SeriesDescription
 
 
 def is_attenuation_corrected(dcm: DICOM) -> bool:
@@ -33,14 +33,14 @@ def is_attenuation_corrected(dcm: DICOM) -> bool:
 
 
 def shows_segments(dcm: DICOM) -> bool:
-    return "PolarMapSeg" in dcm.SeriesDescription
+    return "segment" in dcm.SeriesDescription.lower()
 
 
 def get_type(dcm: DICOM) -> str:
     description = dcm.SeriesDescription.lower()
-    if "syn" in description:
-        return "synthetic"
-    elif "ct" in description:
+    if "DLAC" in description:
+        return "dl"
+    elif "CTAC" in description:
         return "ct"
     else:
         return "symbia"
@@ -100,10 +100,17 @@ if __name__ == "__main__":
         help="the csv file containing meta information of the dataset of which the generated polar maps are",
     )
     parser.add_argument(
-        "--rect",
+        "--rect_upper",
         type=int,
         nargs=4,
-        default=[19, 938, 649, 1568],
+        default=[22, 1499, 22 + 588, 1500 + 588],
+        help="rectangle as [top, left, bottom, right] coordinates where the polar map is found in the DICOM images",
+    )
+    parser.add_argument(
+        "--rect_lower",
+        type=int,
+        nargs=4,
+        default=[634, 1499, 634 + 588, 1500 + 588],
         help="rectangle as [top, left, bottom, right] coordinates where the polar map is found in the DICOM images",
     )
     parser.add_argument(
@@ -143,6 +150,8 @@ if __name__ == "__main__":
         os.mkdir(args.images_dir)
 
     meta = pd.read_csv(args.meta_csv)
+    meta[meta_headers.datetime_acquisition] = pd.to_datetime(meta[meta_headers.datetime_acquisition])
+
     dcm_files = sorted(get_files_recursive(args.polar_map_dir))
 
     data = pd.DataFrame(
@@ -168,78 +177,56 @@ if __name__ == "__main__":
 
         protocol = get_protocol(dcm)
         patient_id = int(dcm.PatientID)
-        meta_row = meta[
+        meta_rows = meta[
             (meta[meta_headers.patient_id] == patient_id)
             & (meta[meta_headers.protocol] == protocol)
-        ].iloc[0]
+            & (meta[meta_headers.datetime_acquisition] == DICOMTime.Acquisition.to_datetime(dcm))
+        ]
+        assert len(meta_rows) == 1
+        meta_row = meta_rows.iloc[0]
 
         row = {
             headers.id: meta_row[meta_headers.id],
             headers.sc: is_scatter_corrected(dcm),
             headers.ac: is_attenuation_corrected(dcm),
             headers.segments: shows_segments(dcm),
-            headers.type: get_type(dcm),
         }
 
-        top, left, bottom, right = args.rect
-        polar_map = img[top:bottom, left:right]
-        polar_map = cv.cvtColor(polar_map, cv.COLOR_RGB2BGR)
-
-        # Note: the images created by Syngo.Via show polar maps of rest studies
-        # further down per default. This little bit of code selects this area
-        # further down if correcting this was forgotten
-        if polar_map.mean() < 10:
-            offset = 649
-            top = top + offset
-            bottom = bottom + offset
+        types = dcm.SeriesDescription.split("-")[-1]
+        types = types[:-1].strip()
+        types = types.split("/")
 
+        for _type, rect in zip(types, [args.rect_upper, args.rect_lower]):
+            top, left, bottom, right = rect
             polar_map = img[top:bottom, left:right]
             polar_map = cv.cvtColor(polar_map, cv.COLOR_RGB2BGR)
 
-        _ac = "ac" if row[headers.ac] else "nac"
-        _sc = "sc" if row[headers.sc] else "nsc"
-        _seg = "segments" if row[headers.segments] else "gray"
-        _file = f"{row[headers.id]:04d}-{_seg}_{_ac}_{_sc}_{row[headers.type]}.png"
-        row[headers.file] = _file
-        _file = os.path.join(args.images_dir, _file)
-
-        if (
-            len(
-                data[
-                    (data[headers.id] == row[headers.id])
-                    & (data[headers.sc] == row[headers.sc])
-                    & (data[headers.ac] == row[headers.ac])
-                    & (data[headers.segments] == row[headers.segments])
-                    & (data[headers.type] == row[headers.type])
-                ]
-            )
-            > 0
-        ):
-            logger.warning(f"Skip {dcm_file} as it is a duplicate for row: {row}")
-            continue
-
-        logger.debug(f"For series {dcm.SeriesDescription} store row {row}")
-
-        store = not args.interactive
-        if args.interactive:
-            while True:
-                cv.imshow("Polar Map", cv.resize(polar_map, (512, 512)))
-                key = cv.waitKey(0)
-
-                if key == ord("q"):
-                    exit(0)
-                elif key == ord("n"):
-                    break
-                elif key == ord("s"):
-                    store = True
-                    break
-
-        if store:
-            cv.imwrite(_file, polar_map)
-
-            row = pd.DataFrame(row, index=[0])
-            data = pd.concat((data, row), ignore_index=True)
-
-            data = data.sort_values(by=[headers.id, headers.file])
-            data.to_csv(args.csv, index=False)
+            row[headers.type] = _type[:2].lower()
+            _seg = "segments" if row[headers.segments] else "gray"
+            _file = f"{row[headers.id]:04d}-{_seg}_ac_nsc_{row[headers.type]}.png"
+            row[headers.file] = _file
+            _file = os.path.join(args.images_dir, _file)
+
+            logger.debug(f"For series {dcm.SeriesDescription} store row {row}")
+
+            store = not args.interactive
+            if args.interactive:
+                while True:
+                    cv.imshow("Polar Map", cv.resize(polar_map, (1024, 1024)))
+                    key = cv.waitKey(0)
+
+                    if key == ord("q"):
+                        exit(0)
+                    elif key == ord("n"):
+                        break
+                    elif key == ord("s"):
+                        store = True
+                        break
+
+            if store:
+                cv.imwrite(_file, polar_map)
+
+                data = pd.concat((data, pd.DataFrame(row, index=[0])), ignore_index=True)
+                data = data.sort_values(by=[headers.id, headers.file])
+                data.to_csv(args.csv, index=False)