Vis fuse2

fuse/data/ops/ops_visprobe.py

@@ -1,23 +1,67 @@
-from typing import Callable, Dict, List, Optional, OrderedDict, Sequence, Tuple, Union
-import copy
+from typing import List, Optional, Union , Dict , Any , Callable
 import enum
+import numpy as np
 
+from pycocotools import mask as maskUtils
 from fuse.utils.ndict import NDict
-from fuse.data.visualizer.visualizer_base import VisualizerBase
+from fuseimg.utils.visualization.visualizer_base import VisualizerBase
+from fuseimg.utils.typing.key_types_imaging import DataTypeImaging
+from torch import Tensor
 from .op_base import OpBase
 from fuse.data.key_types import TypeDetectorBase
 
 class VisFlag(enum.IntFlag):
-    COLLECT = 1       #save current state for future comparison        
-    SHOW_CURRENT = 2          #show current state                                    
-    SHOW_COLLECTED = 4 #show comparison of all previuosly collected states
-    CLEAR = 8           #clear all collected states until this point in the pipeline
-    ONLINE = 16         #show operations will prompt the user with the releveant plot 
-    OFFLINE = 32        #show operations will write to disk (using the caching mechanism) the relevant info (state or states for comparison) 
-    FORWARD = 64        #visualization operation will be activated on forward pipeline execution flow
-    REVERSE = 128       #visualization operation will be activated on reverse pipeline execution flow
-    SHOW_ALL_COLLECTED = 256 #show comparison of all previuosly collected states
+    COLLECT = 1                     #save current state for future comparison        
+    VISUALIZE_CURRENT = 2           #show current state                                    
+    VISUALIZE_COLLECTED = 4          #show comparison of all previuosly collected states
+    CLEAR = 8                       #clear all collected states until this point in the pipeline
+    ONLINE = 16                     #show operations will prompt the user with the releveant plot 
+    # REVERSE = 32                  #visualization operation will be activated on reverse pipeline execution flow
+
+
+def convert_uncompressed_RLE_COCO_type(element : Dict ,height : int  ,width: int)-> np.ndarray:
+    """
+    converts uncompressed RLE to COCO default type ( compressed RLE)
+    :param element:  input in uncompressed Run Length Encoding (RLE - https://en.wikipedia.org/wiki/Run-length_encoding) 
+                    saved in map object example :  {"size": [333, 500],"counts": [26454, 2, 651, 3, 13, 1]}
+                    counts first element is how many bits are not the object, then how many bits are the object and so on
+    :param height: original image height in pixels
+    :param width: original image width in pixels
+    :return  output mask 
+    """
+    p = maskUtils.frPyObjects(element, height, width)
+    return p 
+
+def convert_compressed_RLE_COCO_type(element : list ,height : int  ,width: int)-> np.ndarray:
+    """
+    converts polygon to COCO default type ( compressed RLE)
+    :param element:  polygon - array of X,Y coordinates saves as X.Y , example: [[486.34, 239.01, 477.88, 244.78]]
+    :param height: original image height in pixels
+    :param width: original image width in pixels
+    :return   output mask 
+    """
+    rles = maskUtils.frPyObjects(element, height, width)
+    p = maskUtils.merge(rles) 
+    p = maskUtils.decode(p)
+    return p  
+
 
+
+def convert_COCO_to_mask(elements : Any,height : int  ,width: int, segmentation_type : DataTypeImaging )-> Dict:
+    """
+    converts COCO type to mask
+    :param elements:  input in any COCO format
+    :param height: original image height in pixels
+    :param width: original image width in pixels
+    :param segmentation_type: DataTypeImaging
+    :return  output mask 
+    """     
+    if segmentation_type == DataTypeImaging.UCRLE:
+        elements = [convert_uncompressed_RLE_COCO_type(element,height,width)  for element in elements]    
+    elif segmentation_type == DataTypeImaging.CRLE:
+        elements = [convert_compressed_RLE_COCO_type(element,height,width)  for element in elements]     
+    return elements
+
 class VisProbe(OpBase):
     """
     Handle visualization, saves, shows and compares the sample with respect to the current state inside a pipeline
@@ -26,44 +70,50 @@ class VisProbe(OpBase):
 
     Important notes:
     - running in a cached environment is dangerous and is prohibited
-    - this Operation is not thread safe ans so multithreading is also discouraged
+    - this Operation is not thread safe and so multithreading is also discouraged
 
     "
     """
 
     def __init__(self,flags: VisFlag, 
                  keys: Union[List, dict] ,
                  type_detector: TypeDetectorBase,
-                 id_filter: Union[None, List] = None, 
+                 coco_converter : Callable = None,
+                 name : str ="",
+                 sample_id: Union[None, List] = None, 
                  visualizer: VisualizerBase = None,
-                 cache_path: str = "~/"):
+                 output_path: str = "~/"):
         """ 
           :param flags: operation flags (or possible concatentation of flags using IntFlag), details:   
               COLLECT - save current state for future comparison                                                                                               
-              SHOW_CURRENT - show current state                                                                                                                        
-              SHOW_COllected - show comparison of all previuosly collected states                                                                              
+              VISUALIZE_CURRENT - visualize current state                                                                                                                        
+              VISUALIZE_COLLECTED - visualize comparison of all previuosly collected states                                                                              
               CLEAR - clear all collected states until this point in the pipeline                                                                              
-              ONLINE - show operations will prompt the user with the releveant plot                                                                        
-              OFFLINE - show operations will write to disk (using the caching mechanism) the relevant info (state or states for comparison)                
+              ONLINE - visualize operations will prompt the user with the releveant plot                                                                        
+              OFFLINE - visualize operations will write to disk (using the caching mechanism) the relevant info (state or states for comparison)                
               FORWARD - visualization operation will be activated on forward pipeline execution flow                                                       
               REVERSE - visualization operation will be activated on reverse pipeline execution flow
           :param keys: for which sample keys to handle visualization, also can be grouped in a dictionary
-          :param id_filter: for which sample id's to be activated, if None, active for all samples  
+          :param type_detector : class used to identify objects from the keys
+          :param coco_converter : function that convert the input format to COCO type format
+          :param name : image name to represnt, if not given a number will be shown.
+          :param sample_id: for which sample id's to be activated, if None, active for all samples  
           :param visualizer: the actual visualization handler, depands on domain and use case, should implement Visualizer Base
-          :param cache_path: root dir to save the visualization outputs in offline mode
+          :param output_path: root dir to save the visualization outputs in offline mode
 
           few issues to be aware of, detailed in github issues regarding static cached pipeline and multiprocessing
           note - if both forward and reverse are on, then by default, on forward we do collect and on reverse we do show_collected to 
           compare reverse operations
           for each domain we inherit for VisProbe like ImagingVisProbe,...
 """
         super().__init__()
-        self._id_filter = id_filter
+        self._sample_id = sample_id
         self._keys  = keys
         self._flags = flags
-        self._cacher = None
+        self._coco_converter = coco_converter
+        self._name = name
         self._collected_prefix = "data.$vis"
-        self._cache_path = cache_path
+        self._output_path = output_path
         self._visualizer = visualizer
         self._type_detector = type_detector
 
@@ -76,111 +126,68 @@ def _extract_collected(self, sample_dict: NDict):
                 res.append(vdata)
         return res
 
-    def _extract_data(self, sample_dict: NDict, keys, op_id):
-        if type(keys) is list:
-            # infer keys groups
-            keys.sort()
-            first_type = self._type_detector.get_type(sample_dict, keys[0])
-            num_of_groups = len([self._type_detector.get_type(sample_dict, k) for k in keys if self._type_detector.get_type(sample_dict, k) == first_type])
-            keys_per_group = len(keys) // num_of_groups
-            keys = {f"group{i}": keys[i:i + keys_per_group] for i in range(0, len(keys), keys_per_group)}
-
+    def _extract_data(self, sample_dict: NDict, keys : List , name : str):
         res = NDict()
-        for group_id, group_keys in keys.items():
-            for key in group_keys:
-                prekey = f'groups.{group_id}.{key.replace(".", "_")}'
-                res[f'{prekey}.value'] = sample_dict[key]
-                res[f'{prekey}.type'] = self._type_detector.get_type(sample_dict, key)
-        res['$step_id'] = op_id
+        for key in keys:
+            prekey = key.replace(".", "_")
+            if isinstance(sample_dict[key] , Tensor ):
+                res[f'{prekey}.value'] = sample_dict[key].clone()
+            else :
+                 res[f'{prekey}.value'] = sample_dict[key].copy()
+            res[f'{prekey}.type'] = self._type_detector.get_type(sample_dict, key)
+
+            if self._coco_converter != None :
+                 res[f'{prekey}.value'] = self._coco_converter(res[f'{prekey}.value'],res[f'{prekey}.type'])
+            if res[f'{prekey}.type'] in [DataTypeImaging.CRLE, DataTypeImaging.UCRLE  ]:
+                res[f'{prekey}.value'] = convert_COCO_to_mask(res[f'{prekey}.value'],sample_dict['height'],sample_dict['width'],res[f'{prekey}.type'] )
+                res[f'{prekey}.type'] = DataTypeImaging.SEG 
+            res[f'{prekey}.name'] = name
         return res
 
-
-    def _save(self, vis_data: Union[List, dict]):
-        # use caching to save all relevant vis_data
-        print("saving vis_data", vis_data)
-
-    def _handle_flags(self, flow, sample_dict: NDict, op_id: Optional[str]):
+    def _handle_flags(self, sample_dict: NDict, op_id: Optional[str]):
         """
         See super class
         """
         # sample was filtered out by its id
-        if self._id_filter and self.get_idx(sample_dict) not in self._id_filter:
+        if self._sample_id and self.get_idx(sample_dict) not in self._sample_id:
             return None
-        if flow not in self._flags:
-            return None 
 
-        # grouped key dictionary with the following structure:
-        #vis_data = {"cc_group":
-        #                {
-        #                    "key1": {
-        #                      "value": ndarray, 
-        #                      "type": DataType.Image,
-        #                      "op_id": "test1"}
-        #                    "key2": {
-        #                      "value": ndarray, 
-        #                      "type": DataType.BBox,
-        #                      "op_id": "test1"}
-        #                 },
-        #            "mlo_goup":
-        #                {
-        #                    "key3": {
-        #                      "value": ndarray, 
-        #                      "type": DataType.Image,
-        #                      "op_id": "test1"}
-        #                    "key4": {
-        #                      "value": ndarray, 
-        #                      "type": DataType.BBox,
-        #                      "op_id": "test1"}
-        #                 },
-        #            }
-        vis_data = self._extract_data(sample_dict, self._keys, op_id)
-        both_fr =  (VisFlag.REVERSE | VisFlag.FORWARD) in self._flags
-        dir_forward = flow == VisFlag.FORWARD
-        dir_reverse = flow == VisFlag.REVERSE
-        any_show_collected = VisFlag.SHOW_ALL_COLLECTED|VisFlag.SHOW_COLLECTED
-
-        if VisFlag.COLLECT in self._flags or (dir_forward and both_fr):
+        vis_data = self._extract_data(sample_dict, self._keys ,self._name)
+        name_prefix=""
+        if "name" in sample_dict.to_dict().keys():
+            name_prefix = sample_dict["name"]
+        name_prefix += "."+op_id
+        if VisFlag.COLLECT in self._flags:
             if not self._collected_prefix in sample_dict:
                 sample_dict[self._collected_prefix] = []
             sample_dict[self._collected_prefix].append(vis_data)
 
 
-        if VisFlag.SHOW_CURRENT in self._flags:
+        if VisFlag.VISUALIZE_CURRENT in self._flags:
             if VisFlag.ONLINE in self._flags:
                 self._visualizer.show(vis_data)
-            if VisFlag.OFFLINE in self._flags:
-                self._save(vis_data)
+            else :
+                self._visualizer.save(vis_data , name_prefix , self._output_path)
 
-        if (VisFlag.SHOW_ALL_COLLECTED in self._flags or VisFlag.SHOW_COLLECTED in self._flags) and (
-            (both_fr and dir_reverse) or not both_fr):
+        if  VisFlag.VISUALIZE_COLLECTED in self._flags:
             vis_data = self._extract_collected(sample_dict) + [vis_data]
-            if both_fr:
-                if VisFlag.SHOW_COLLECTED in self._flags:
-                    vis_data = vis_data[-2:]
             if VisFlag.ONLINE in self._flags:
                 self._visualizer.show(vis_data)
-            if VisFlag.OFFLINE in self._flags:
-                self.save(vis_data)
+            else :
+                self._visualizer.save(vis_data , name_prefix, self._output_path)
 
         if VisFlag.CLEAR in self._flags:
             sample_dict[self._collected_prefix] = []
 
-        if VisFlag.SHOW_COLLECTED in self._flags and both_fr and dir_reverse:
-            sample_dict[self._collected_prefix].pop()
+        # TODO - support reverse?
+        # if VisFlag.SHOW_COLLECTED in self._flags and VisFlag.REVERSE in self._flags:
+        #     sample_dict[self._collected_prefix].pop()
 
         return sample_dict
 
 
     def __call__(self, sample_dict: NDict, op_id: Optional[str], **kwargs) -> Union[None, dict, List[dict]]:
-        res = self._handle_flags(VisFlag.FORWARD, sample_dict, op_id)
+        res = self._handle_flags(sample_dict, op_id)
         return res
 
-    def reverse(self, sample_dict: NDict, op_id: Optional[str], key_to_reverse: str, key_to_follow: str) -> dict:
-        """
-        See super class
-        """
-        res = self._handle_flags(VisFlag.REVERSE, sample_dict, op_id)
-        if res is None:
-            res = sample_dict
-        return res