Merge pull request #50 from WEHI-ResearchComputing/45-mrc-compatibility-when-not-using-denoiser

Add MRC Input Compatibility When Not Using Denoiser

Closes #45

TO-DO
Review raw MRC detection weights to ensure optimal performance when not using denoiser: https://huggingface.co/MihinP/PartiNet

Author: MihinP

Singularity

@@ -18,4 +18,4 @@ from: python:3.9.19-slim-bookworm
 %labels
   AUTHORS Mihin Perera, Edward Yang, Julie Iskander
   MAINTAINERS Mihin Perera, Edward Yang, Julie Iskander
-  VERSION v0.2.0
+  VERSION v1.0.0

partinet/DynamicDet/detect.py

@@ -41,6 +41,17 @@
 sys.stdout.reconfigure(line_buffering=True)
 sys.stderr.reconfigure(line_buffering=True)
 
+# helper ----------------------------------------------------------------
+def _save_filename(p: Path) -> str:
+    """Return an output filename for image `p`.
+
+    Currently we cannot write `.mrc` files with OpenCV, so convert any
+    input with that suffix to PNG.  Other extensions are returned as-is.
+    """
+    if p.suffix.lower() == '.mrc':
+        return p.stem + '.png'
+    return p.name
+
 # --- Detection function ---
 def detect(opt, save_img=False):
     source, cfg, weight, view_img, save_txt, nc, imgsz = (
@@ -145,7 +156,8 @@ def worker(device):
                     p, s, im0, f_idx = path, '', im0s, frame
 
                 p = Path(p)
-                save_path = str(save_dir / p.name)
+                # choose image-filename for saving; helper contains mrc logic
+                save_path = str(save_dir / _save_filename(p))
                 txt_path = str(save_dir / 'labels' / p.stem) + ('' if dataset.mode == 'image' else f'_{f_idx}')
                 gn = torch.tensor(im0.shape)[[1, 0, 1, 0]]
 

partinet/DynamicDet/test.py

@@ -13,7 +13,7 @@
 
 import partinet.DynamicDet
 from partinet.DynamicDet.models.yolo import Model
-from partinet.DynamicDet.utils.datasets import create_dataloader
+from partinet.DynamicDet.utils.datasets import create_dataloader, LoadImages
 from partinet.DynamicDet.utils.general import coco80_to_coco91_class, check_dataset, check_file, check_img_size, \
     box_iou, non_max_suppression, scale_coords, xyxy2xywh, xywh2xyxy, set_logging, increment_path, colorstr
 from partinet.DynamicDet.utils.metrics import ap_per_class, ConfusionMatrix

partinet/DynamicDet/utils/datasets.py

@@ -19,6 +19,7 @@
 from PIL import Image, ExifTags
 from torch.utils.data import Dataset
 from tqdm import tqdm
+import mrcfile  # support cryoEM micrograph format
 
 import pickle
 from copy import deepcopy
@@ -29,10 +30,14 @@
 from partinet.DynamicDet.utils.general import check_requirements, xyxy2xywh, xywh2xyxy, xywhn2xyxy, xyn2xy, segment2box, segments2boxes, \
     resample_segments, clean_str
 from partinet.DynamicDet.utils.torch_utils import torch_distributed_zero_first
+from partinet.process_utils.guided_denoiser import transform
 
 # Parameters
 help_url = 'https://github.com/ultralytics/yolov5/wiki/Train-Custom-Data'
-img_formats = ['bmp', 'jpg', 'jpeg', 'png', 'tif', 'tiff', 'dng', 'webp', 'mpo']  # acceptable image suffixes
+# acceptable image suffixes (extensions are compared case‑insensitively)
+# add 'mrc' so that LoadImages/LoadImagesAndLabels can see micrographs directly
+# micrographs are read with `mrcfile`; only uncompressed MRC files are supported
+img_formats = ['bmp', 'jpg', 'jpeg', 'png', 'tif', 'tiff', 'dng', 'webp', 'mpo', 'mrc']  
 vid_formats = ['mov', 'avi', 'mp4', 'mpg', 'mpeg', 'm4v', 'wmv', 'mkv']  # acceptable video suffixes
 logger = logging.getLogger(__name__)
 
@@ -137,6 +142,7 @@ def __init__(self, path, img_size=640, stride=32):
         else:
             raise Exception(f'ERROR: {p} does not exist')
 
+        # allow uncompressed .mrc files
         images = [x for x in files if x.split('.')[-1].lower() in img_formats]
         videos = [x for x in files if x.split('.')[-1].lower() in vid_formats]
         ni, nv = len(images), len(videos)
@@ -183,7 +189,15 @@ def __next__(self):
         else:
             # Read image
             self.count += 1
-            img0 = cv2.imread(path)  # BGR
+            # support cryo-EM micrographs saved as MRC
+            if path.lower().endswith('.mrc'):
+                img_mrc = mrcfile.read(path)
+                img0 = transform(img_mrc).astype(np.uint8)
+                # `transform` returns a single-channel image; networks expect BGR
+                if img0.ndim == 2:
+                    img0 = cv2.cvtColor(img0, cv2.COLOR_GRAY2BGR)
+            else:
+                img0 = cv2.imread(path)  # BGR
             assert img0 is not None, 'Image Not Found ' + path
             #print(f'image {self.count}/{self.nf} {path}: ', end='')
 
@@ -668,7 +682,14 @@ def load_image(self, index):
     img = self.imgs[index]
     if img is None:  # not cached
         path = self.img_files[index]
-        img = cv2.imread(path)  # BGR
+        # support uncompressed mrc micrographs
+        if path.lower().endswith('.mrc'):
+            img_mrc = mrcfile.read(path)
+            img = transform(img_mrc).astype(np.uint8)
+            if img.ndim == 2:
+                img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)
+        else:
+            img = cv2.imread(path)  # BGR
         assert img is not None, 'Image Not Found ' + path
         h0, w0 = img.shape[:2]  # orig hw
         r = self.img_size / max(h0, w0)  # resize image to img_size

partinet/__init__.py

@@ -1,7 +1,7 @@
 import click
 import sys, os
 
-__version__ = "0.2.0"
+__version__ = "1.0.0"
 
 DYNAMICDET_AVAILABLE_MODELS = ["yolov7", "yolov7x", "yolov7-w6", "yolov7-e6", "yolov7-d6", "yolov7-e6e"]