a

Author: armarpc

mmseg/datasets/pipelines/hubmap_transforms.py

@@ -68,6 +68,29 @@ def __call__(self, results):
 
         return results
 
+
+@PIPELINES.register_module()
+class HubmapTestTransform:
+
+    def __call__(self, results):
+        img = results['img']
+        keys = []
+        masks = []
+        for key in results.get('seg_fields', []):
+            keys.append(key)
+            masks.append(results[key])
+        
+        img = np.float32(img) if img.dtype != np.float32 else img.copy()
+        assert img.dtype != np.uint8
+        #if to_rgb:
+        cv2.cvtColor(img, cv2.COLOR_BGR2RGB, img)  # inplace
+
+        results['img'] = img
+        for key, mask in zip(keys, masks):
+            results[key] = mask
+
+        return results
+
 #from albumentations.pytorch import ToTensor
 #from get_config import get_config