(catalyst 20.03): update code

Author: bagxi

bin/tests/check_instance.sh

@@ -49,7 +49,7 @@ print(iou_soft)
 print(iou_hard)
 
 assert aggregated_loss < 0.9
-assert iou_soft > 0.05
+assert iou_soft > 0.04
 assert iou_hard > 0.1
 """
 

configs/templates/binary.yml

@@ -118,9 +118,6 @@ stages:
       callback: RawMaskPostprocessingCallback
       output_key: binary_mask
 
-    raw_processor:
-      callback: RawMaskPostprocessingCallback
-
     iou_soft:
       callback: IouCallback
       input_key: mask

configs/templates/instance.yml

@@ -10,7 +10,7 @@ args:
 stages:
 
   state_params:
-    main_metric: &reduce_metric iou_hard
+    main_metric: &reduced_metric iou_hard
     minimize_metric: False
 
   data_params:
@@ -56,7 +56,7 @@ stages:
               contrast_limit: 0.5
             - transform: A.RandomGamma
             - transform: A.CLAHE
-            - transform: A.JpegCompression
+            - transform: A.ImageCompression
               quality_lower: 50
         - &post_transforms
           transform: A.Compose
@@ -108,10 +108,11 @@ stages:
       multiplier: 1.0
 
     loss_aggregator:
-      callback: CriterionAggregatorCallback
+      callback: MetricAggregationCallback
       prefix: &aggregated_loss loss
-      loss_aggregate_fn: "mean" # or "sum"
-      multiplier: 1.0 # scale factor for the aggregated loss
+      metrics: [loss_bce, loss_dice, loss_iou]
+      mode: "mean"
+      multiplier: 1.0
 
     raw_processor:
       callback: RawMaskPostprocessingCallback
@@ -140,7 +141,7 @@ stages:
       loss_key: *aggregated_loss
     scheduler:
       callback: SchedulerCallback
-      reduce_metric: *reduce_metric
+      reduced_metric: *reduced_metric
     saver:
       callback: CheckpointCallback
 

configs/templates/semantic.yml

@@ -118,10 +118,6 @@ stages:
       callback: RawMaskPostprocessingCallback
       output_key: semantic_mask
 
-    raw_processor:
-      callback: RawMaskPostprocessingCallback
-      output_key: semantic_mask
-
     iou_soft:
       callback: IouCallback
       input_key: mask

scripts/process_instance_masks.py

@@ -14,10 +14,7 @@
 
 def build_args(parser):
     parser.add_argument(
-        "--in-dir",
-        type=Path,
-        required=True,
-        help="Raw masks folder path"
+        "--in-dir", type=Path, required=True, help="Raw masks folder path"
     )
     parser.add_argument(
         "--out-dir",
@@ -151,10 +148,10 @@ def preprocess(self, sample: Path):
                     else:
                         sz = 3
 
-                uniq = np.unique(labels[
-                    max(0, y0 - sz):min(labels.shape[0], y0 + sz + 1),
-                    max(0, x0 - sz):min(labels.shape[1], x0 + sz + 1),
-                ])
+                uniq = np.unique(
+                    labels[max(0, y0 - sz):min(labels.shape[0], y0 + sz + 1),
+                           max(0, x0 - sz):min(labels.shape[1], x0 + sz + 1)]
+                )
                 if len(uniq[uniq > 0]) > 1:
                     borders[y0, x0] = 255
                     mask_without_borders[y0, x0] = 0

src/callbacks/io.py

@@ -3,7 +3,7 @@
 import imageio
 import numpy as np
 
-from catalyst.dl import Callback, CallbackOrder, State, utils
+from catalyst.dl import Callback, CallbackNode, CallbackOrder, State, utils
 from .utils import crop_by_masks, mask_to_overlay_image
 
 
@@ -17,7 +17,7 @@ def __init__(
         input_key: str = "image",
         outpath_key: str = "name",
     ):
-        super().__init__(CallbackOrder.Logging)
+        super().__init__(order=CallbackOrder.Logging, node=CallbackNode.Master)
         self.output_dir = Path(output_dir)
         self.relative = relative
         self.filename_suffix = filename_suffix
@@ -102,10 +102,11 @@ def __init__(
         self.output_key = output_key
 
     def on_batch_end(self, state: State):
-        names = state.input[self.outpath_key]
-        images = utils.tensor_to_ndimage(state.input[self.input_key].cpu())
-        masks = state.output[self.output_key]
+        names = state.batch_in[self.outpath_key]
+        images = state.batch_in[self.input_key]
+        masks = state.batch_out[self.output_key]
 
+        images = utils.tensor_to_ndimage(images.detach().cpu())
         for name, image, masks_ in zip(names, images, masks):
             instances = crop_by_masks(image, masks_)
 

src/callbacks/processing.py

@@ -1,6 +1,6 @@
 import torch
 
-from catalyst.dl import Callback, CallbackOrder, State
+from catalyst.dl import Callback, CallbackNode, CallbackOrder, State
 from .utils import encode_mask_with_color, label_instances
 
 
@@ -11,7 +11,7 @@ def __init__(
         input_key: str = "logits",
         output_key: str = "mask",
     ):
-        super().__init__(CallbackOrder.Internal)
+        super().__init__(order=CallbackOrder.Internal, node=CallbackNode.All)
         self.threshold = threshold
         self.input_key = input_key
         self.output_key = output_key
@@ -21,7 +21,7 @@ def on_batch_end(self, state: State):
 
         output = torch.sigmoid(output).detach().cpu().numpy()
         state.batch_out[self.output_key] = encode_mask_with_color(
-            output, self.threshold
+            output, threshold=self.threshold
         )
 
 
@@ -35,7 +35,7 @@ def __init__(
         out_key_semantic: str = None,
         out_key_border: str = None,
     ):
-        super().__init__(CallbackOrder.Internal)
+        super().__init__(CallbackOrder.Internal, node=CallbackNode.All)
         self.watershed_threshold = watershed_threshold
         self.mask_threshold = mask_threshold
         self.input_key = input_key
@@ -44,22 +44,22 @@ def __init__(
         self.out_key_border = out_key_border
 
     def on_batch_end(self, state: State):
-        output: torch.Tensor = torch.sigmoid(state.output[self.input_key])
+        output = state.batch_out[self.input_key]
 
+        output = torch.sigmoid(output).detach().cpu()
         semantic, border = output.chunk(2, -3)
 
         if self.out_key_semantic is not None:
-            state.output[self.out_key_semantic] = encode_mask_with_color(
-                semantic.data.cpu().numpy(), threshold=self.mask_threshold
+            state.batch_out[self.out_key_semantic] = encode_mask_with_color(
+                semantic.numpy(), threshold=self.mask_threshold
             )
 
         if self.out_key_border is not None:
-            state.output[self.out_key_border] = (
-                border.data.cpu().squeeze(-3).numpy() >
-                self.watershed_threshold
+            state.batch_out[self.out_key_border] = (
+                border.squeeze(-3).numpy() > self.watershed_threshold
             )
 
-        state.output[self.output_key] = label_instances(
+        state.batch_out[self.output_key] = label_instances(
             semantic,
             border,
             watershed_threshold=self.watershed_threshold,