deepliif 1.2.3__py3-none-any.whl → 1.2.4__py3-none-any.whl

deepliif/models/__init__.py

@@ -133,7 +133,7 @@ def load_eager_models(opt, devices=None):
         model_names = list(devices.keys())
     else:
         model_names = model.model_names
-    
+
     for name in model_names:#model.model_names:
         if isinstance(name, str):
             if '_' in name:
@@ -155,7 +155,6 @@ def load_eager_models(opt, devices=None):
             
     return nets
 
-
 @lru_cache
 def init_nets(model_dir, eager_mode=False, opt=None, phase='test'):
     """
@@ -169,14 +168,20 @@ def init_nets(model_dir, eager_mode=False, opt=None, phase='test'):
         opt = get_opt(model_dir, mode=phase)
     opt.use_dp = False
     
+    # get l_net_groups: a list of lists, each sublist is a candidate, ordered by priority (try 1st, if fail try 2nd, etc.)
     if opt.model in ['DeepLIIF','DeepLIIFKD']:
-        net_groups = [
-            ('G1', 'G52'),
-            ('G2', 'G53'),
-            ('G3', 'G54'),
-            ('G4', 'G55'),
-            ('G51',)
-        ]
+        # net_groups = [
+        #     ('G1', 'G52'),
+        #     ('G2', 'G53'),
+        #     ('G3', 'G54'),
+        #     ('G4', 'G55'),
+        #     ('G51',)
+        # ]
+        if opt.modalities_no == 0:
+            net_groups = [(f'G{opt.mod_id_seg}{opt.input_id}',)]
+        else:
+            net_groups = [(f'G{i+1}', f'G{opt.mod_id_seg}{int(opt.input_id)+i+1}') for i in range(opt.modalities_no)]
+            net_groups += [(f'G{opt.mod_id_seg}{opt.input_id}',)] # this is the generator for the input base mod
     elif opt.model in ['DeepLIIFExt','SDG']:
         if opt.seg_gen:
             net_groups = [(f'G_{i+1}',f'GS_{i+1}') for i in range(opt.modalities_no)]
@@ -197,8 +202,8 @@ def init_nets(model_dir, eager_mode=False, opt=None, phase='test'):
         mapping_gpu_ids = {i:idx for i,idx in enumerate(opt.gpu_ids)}
         chunks = [itertools.chain.from_iterable(c) for c in chunker(net_groups, number_of_gpus)]
         # chunks = chunks[1:]
-        devices = {n: torch.device(f'cuda:{mapping_gpu_ids[i]}') for i, g in enumerate(chunks) for n in g}
-        # devices = {n: torch.device(f'cuda:{i}') for i, g in enumerate(chunks) for n in g}
+        #l_devices = [{n: torch.device(f'cuda:{mapping_gpu_ids[i]}') for i, g in enumerate(chunks) for n in g} for chunks in l_chunks]
+        devices = {n: torch.device(f'cuda:{i}') for i, g in enumerate(chunks) for n in g}
     else:
         devices = {n: torch.device('cpu') for n in itertools.chain.from_iterable(net_groups)}
 
@@ -219,13 +224,14 @@ def compute_overlap(img_size, tile_size):
     return tile_size // 4
 
 
-def run_torchserve(img, model_path=None, nets=None, eager_mode=False, opt=None, seg_only=False, seg_weights=None, use_dask=True, output_tensor=False):
+def run_torchserve(img, model_path=None, nets=None, eager_mode=False, opt=None, seg_only=False, mod_only=False, seg_weights=None, use_dask=True, output_tensor=False):
     """
     eager_mode: not used in this function; put in place to be consistent with run_dask
            so that run_wrapper() could call either this function or run_dask with
            same syntax
     opt: same as eager_mode
     seg_only: same as eager_mode
+    mod_only: same as eager_mode
     seg_weights: same as eager_mode
     nets: same as eager_mode
     """
@@ -246,7 +252,8 @@ def run_torchserve(img, model_path=None, nets=None, eager_mode=False, opt=None,
     return {k: tensor_to_pil(deserialize_tensor(v)) for k, v in res.json().items()}
 
 
-def run_dask(img, model_path=None, nets=None, eager_mode=False, opt=None, seg_only=False, seg_weights=None, use_dask=True, output_tensor=False):
+def run_dask(img, model_path=None, nets=None, eager_mode=False, opt=None, seg_only=False, mod_only=False,
+             seg_weights=None, use_dask=True, output_tensor=False):
     """
     Provide either the model path or the networks object.
     
@@ -269,65 +276,81 @@ def run_dask(img, model_path=None, nets=None, eager_mode=False, opt=None, seg_on
         else:
             ts = transform(img.resize((opt.scale_size, opt.scale_size)))
     
-
+    
     if use_dask:
         @delayed
         def forward(input, model):
             with torch.no_grad():
                 return model(input.to(next(model.parameters()).device))
-    else: # some train settings like spectral norm some how in inference mode is not compatible with dask
+    else: # some train settings like spectral norm somehow in inference mode is not compatible with dask
         def forward(input, model):
             with torch.no_grad():
                 return model(input.to(next(model.parameters()).device))
     
     if opt.model in ['DeepLIIF','DeepLIIFKD']:
         if seg_weights is None:
-            weights = {
-                'G51': 0.5, # IHC
-                'G52': 0.0, # Hema
-                'G53': 0.0, # DAPI
-                'G54': 0.0, # Lap2
-                'G55': 0.5, # Marker
-            }
+            # weights = {
+            #     'G51': 0.5, # IHC
+            #     'G52': 0.0, # Hema
+            #     'G53': 0.0, # DAPI
+            #     'G54': 0.0, # Lap2
+            #     'G55': 0.5, # Marker
+            # }
+            weights = {f'G{opt.mod_id_seg}{int(opt.input_id)+i}': 1/(opt.modalities_no+1) for i in range(opt.modalities_no+1)}
         else:
-            weights = {
-                'G51': seg_weights[0], # IHC
-                'G52': seg_weights[1], # Hema
-                'G53': seg_weights[2], # DAPI
-                'G54': seg_weights[3], # Lap2
-                'G55': seg_weights[4], # Marker
-            }
-
-        seg_map = {'G1': 'G52', 'G2': 'G53', 'G3': 'G54', 'G4': 'G55'}
+            # weights = {
+            #     'G51': seg_weights[0], # IHC
+            #     'G52': seg_weights[1], # Hema
+            #     'G53': seg_weights[2], # DAPI
+            #     'G54': seg_weights[3], # Lap2
+            #     'G55': seg_weights[4], # Marker
+            # }
+            weights = {f'G{opt.mod_id_seg}{int(opt.input_id)+i}': seg_weight for i,seg_weight in enumerate(seg_weights)}
+        
+        
+        seg_map = {f'G{i+1}': f'G{opt.mod_id_seg}{int(opt.input_id)+i+1}' for i in range(opt.modalities_no)}
         if seg_only:
             seg_map = {k: v for k, v in seg_map.items() if weights[v] != 0}
         
         lazy_gens = {k: forward(ts, nets[k]) for k in seg_map}
-        if 'G4' not in seg_map:
-            lazy_gens['G4'] = forward(ts, nets['G4'])
+        if 'Marker' in opt.modalities_names:
+            mod_id_marker = opt.modalities_names.index("Marker")
+            if f'G{mod_id_marker}' not in seg_map:
+                lazy_gens[f'G{mod_id_marker}'] = forward(ts, nets[f'G{mod_id_marker}'])
+        
         gens = compute(lazy_gens)[0]
         
-        lazy_segs = {v: forward(gens[k], nets[v]) for k, v in seg_map.items()}
-        if not seg_only or weights['G51'] != 0:
-            lazy_segs['G51'] = forward(ts, nets['G51'])
-        segs = compute(lazy_segs)[0]
+        if not mod_only:
+            lazy_segs = {v: forward(gens[k], nets[v]) for k, v in seg_map.items()}
+            # run seg generator for the base input
+            if weights[f'G{opt.mod_id_seg}{opt.input_id}'] != 0:
+                lazy_segs[f'G{opt.mod_id_seg}{opt.input_id}'] = forward(ts, nets[f'G{opt.mod_id_seg}{opt.input_id}'])
+            segs = compute(lazy_segs)[0]
         
-        device = next(nets['G1'].parameters()).device # take the device of the first net and move all outputs there for seg aggregation
-        seg = torch.stack([torch.mul(segs[k].to(device), weights[k]) for k in segs.keys()]).sum(dim=0)
+            model_name_first = list(nets.keys())[0]
+            device = next(nets[model_name_first].parameters()).device # take the device of the first net and move all outputs there for seg aggregation
+            seg = torch.stack([torch.mul(segs[k].to(device), weights[k]) for k in segs.keys()]).sum(dim=0)
         
         if output_tensor:
-            if seg_only:
-                res = {'G4': gens['G4']} if 'G4' in gens else {}
+            if mod_only:
+                res = gens
+            elif seg_only and opt.modalities_no > 0:
+                res = {f'G{opt.modalities_no}': gens[f'G{opt.modalities_no}']} if f'G{opt.modalities_no}' in gens else {}
+                res[f'G{opt.mod_id_seg}'] = seg
             else:
                 res = {**gens, **segs}
-            res['G5'] = seg
+                res[f'G{opt.mod_id_seg}'] = seg
+            
         else:
-            if seg_only:
-                res = {'G4': tensor_to_pil(gens['G4'].to(torch.device('cpu')))} if 'G4' in gens else {}
+            if mod_only:
+                res = {k: tensor_to_pil(v.to(torch.device('cpu'))) for k, v in gens.items()}
+            elif seg_only and opt.modalities_no > 0:
+                res = {f'G{opt.modalities_no}': tensor_to_pil(gens[f'G{opt.modalities_no}'].to(torch.device('cpu')))} if f'G{opt.modalities_no}' in gens else {}
+                res[f'G{opt.mod_id_seg}'] = tensor_to_pil(seg.to(torch.device('cpu')))
             else:
                 res = {k: tensor_to_pil(v.to(torch.device('cpu'))) for k, v in gens.items()}
                 res.update({k: tensor_to_pil(v.to(torch.device('cpu'))) for k, v in segs.items()})
-            res['G5'] = tensor_to_pil(seg.to(torch.device('cpu')))
+                res[f'G{opt.mod_id_seg}'] = tensor_to_pil(seg.to(torch.device('cpu')))
     
         return res
     elif opt.model in ['DeepLIIFExt','SDG','CycleGAN']:
@@ -343,6 +366,8 @@ def run_dask(img, model_path=None, nets=None, eager_mode=False, opt=None, seg_on
             gens = {k: forward(ts, nets[k]) for k in seg_map}
         
         res = {k: tensor_to_pil(v) for k, v in gens.items()}
+        if mod_only:
+            return res
     
         if opt.seg_gen:
             if use_dask:
@@ -358,36 +383,55 @@ def run_dask(img, model_path=None, nets=None, eager_mode=False, opt=None, seg_on
 
 
 def is_empty(tile):
-    thresh = 15
+    thresh = 9
     if isinstance(tile, list): # for pair of tiles, only mark it as empty / no need for prediction if ALL tiles are empty
         return all([True if image_variance_gray(t) < thresh else False for t in tile])
     else:
         return True if image_variance_gray(tile) < thresh else False
 
 
-def run_wrapper(tile, run_fn, model_path=None, nets=None, eager_mode=False, opt=None, seg_only=False, seg_weights=None, use_dask=True, output_tensor=False):
+def run_wrapper(tile, run_fn, model_path=None, nets=None, eager_mode=False, opt=None, seg_only=False, mod_only=False, seg_weights=None, use_dask=True, output_tensor=False):
     if opt.model in ['DeepLIIF','DeepLIIFKD']:
         if is_empty(tile):
-            if seg_only:
-                return {
-                    'G4': Image.new(mode='RGB', size=(512, 512), color=(10, 10, 10)),
-                    'G5': Image.new(mode='RGB', size=(512, 512), color=(0, 0, 0)),
+            if seg_only: # return seg image and the last translated modality
+                res = {
+                    #f'G{opt.modalities_no}': Image.new(mode='RGB', size=(512, 512), color=(10, 10, 10)),
+                    f'G{opt.modalities_no}': Image.new(mode='RGB', size=(512, 512), color=opt.background_colors[-1]),
+                    f'G{opt.mod_id_seg}': Image.new(mode='RGB', size=(512, 512), color=(0, 0, 0)),
                 }
+            elif mod_only:
+                res = {f'G{i+1}': Image.new(mode='RGB', size=(512, 512), color=opt.background_colors[i]) for i in range(opt.modalities_no)}
+                
             else :
-                return {
-                    'G1': Image.new(mode='RGB', size=(512, 512), color=(201, 211, 208)),
-                    'G2': Image.new(mode='RGB', size=(512, 512), color=(10, 10, 10)),
-                    'G3': Image.new(mode='RGB', size=(512, 512), color=(0, 0, 0)),
-                    'G4': Image.new(mode='RGB', size=(512, 512), color=(10, 10, 10)),
-                    'G5': Image.new(mode='RGB', size=(512, 512), color=(0, 0, 0)),
-                    'G51': Image.new(mode='RGB', size=(512, 512), color=(0, 0, 0)),
-                    'G52': Image.new(mode='RGB', size=(512, 512), color=(0, 0, 0)),
-                    'G53': Image.new(mode='RGB', size=(512, 512), color=(0, 0, 0)),
-                    'G54': Image.new(mode='RGB', size=(512, 512), color=(0, 0, 0)),
-                    'G55': Image.new(mode='RGB', size=(512, 512), color=(0, 0, 0)),
-                }
+                # return {
+                #     'G1': Image.new(mode='RGB', size=(512, 512), color=(201, 211, 208)),
+                #     'G2': Image.new(mode='RGB', size=(512, 512), color=(10, 10, 10)),
+                #     'G3': Image.new(mode='RGB', size=(512, 512), color=(0, 0, 0)),
+                #     'G4': Image.new(mode='RGB', size=(512, 512), color=(10, 10, 10)),
+                #     'G5': Image.new(mode='RGB', size=(512, 512), color=(0, 0, 0)),
+                #     'G51': Image.new(mode='RGB', size=(512, 512), color=(0, 0, 0)),
+                #     'G52': Image.new(mode='RGB', size=(512, 512), color=(0, 0, 0)),
+                #     'G53': Image.new(mode='RGB', size=(512, 512), color=(0, 0, 0)),
+                #     'G54': Image.new(mode='RGB', size=(512, 512), color=(0, 0, 0)),
+                #     'G55': Image.new(mode='RGB', size=(512, 512), color=(0, 0, 0)),
+                # }
+                res = {**{f'G{i+1}': Image.new(mode='RGB', size=(512, 512), color=opt.background_colors[i]) for i in range(opt.modalities_no)},
+                       **{f'G{opt.mod_id_seg}': Image.new(mode='RGB', size=(512, 512), color=(0, 0, 0))}}
+                
+                # assign mod-wise seg output to corresponding keys, again currently input_id only supports 0 or 1
+                if opt.input_id == 1:
+                    res = {**res,
+                           **{f'G{opt.mod_id_seg}{i+1}': Image.new(mode='RGB', size=(512, 512), color=(0, 0, 0)) for i in range(opt.modalities_no+1)}}
+                else:
+                    res = {**res,
+                           **{f'G{opt.mod_id_seg}{i}': Image.new(mode='RGB', size=(512, 512), color=(0, 0, 0)) for i in range(opt.modalities_no+1)}}
+
+            # when modalities_no = 0... we do not need to generate G0 - output should be just a seg image
+            if 'G0' in res:
+                del res['G0']
+            return res
         else:
-            return run_fn(tile, model_path, None, eager_mode, opt, seg_only, seg_weights)
+            return run_fn(tile, model_path, None, eager_mode, opt, seg_only, mod_only, seg_weights)
     elif opt.model in ['DeepLIIFExt', 'SDG']:
         if is_empty(tile):
             res = {'G_' + str(i): Image.new(mode='RGB', size=(512, 512)) for i in range(1, opt.modalities_no + 1)}
@@ -408,7 +452,8 @@ def run_wrapper(tile, run_fn, model_path=None, nets=None, eager_mode=False, opt=
 
 def inference(img, tile_size, overlap_size, model_path, use_torchserve=False,
               eager_mode=False, color_dapi=False, color_marker=False, opt=None,
-              return_seg_intermediate=False, seg_only=False, seg_weights=None, opt_args={}):
+              return_seg_intermediate=False, seg_only=False, mod_only=False,
+              seg_weights=None, opt_args={}):
     """
     opt_args: a dictionary of key and values to add/overwrite to opt
     """
@@ -433,41 +478,75 @@ def inference(img, tile_size, overlap_size, model_path, use_torchserve=False,
 
     tiler = InferenceTiler(orig, tile_size, overlap_size)
     for tile in tiler:
-        tiler.stitch(run_wrapper(tile, run_fn, model_path, None, eager_mode, opt, seg_only, seg_weights))
+        tiler.stitch(run_wrapper(tile, run_fn, model_path, None, eager_mode, opt, seg_only, mod_only, seg_weights))
         
     results = tiler.results()
-
+    
     if opt.model in ['DeepLIIF','DeepLIIFKD']:
+        # check if both the elements and the order are exactly the same
+        l_modname = [f'mod{i+1}' for i in range(opt.modalities_no)]
+        if l_modname != opt.modalities_names[1:]: 
+            # if not, append modalities_names to mod names
+            l_modname = [f'mod{i+1}-{mod_name}' for i,mod_name in enumerate(opt.modalities_names[1:])]
+        d_modname2id = {mod_name:f'G{i+1}' for i,mod_name in enumerate(l_modname)}
+        
+        if opt.seg_gen:
+            l_modname_seg = [f'mod{i}' for i in range(opt.modalities_no+1)]
+            if l_modname_seg != opt.modalities_names: 
+                # if not, append modalities_names to mod names
+                l_modname_seg = [f'mod{i}-{mod_name}' for i,mod_name in enumerate(opt.modalities_names)]
+            if f'G{opt.mod_id_seg}0' in results.keys():
+                d_modname2id_seg = {mod_name:f'G{opt.mod_id_seg}{i}' for i,mod_name in enumerate(l_modname_seg)}
+            else:
+                d_modname2id_seg = {mod_name:f'G{opt.mod_id_seg}{i+1}' for i,mod_name in enumerate(l_modname_seg)}
+        
+        if not mod_only:
+            d_modname2id['Seg'] = f'G{opt.mod_id_seg}'
+        
+        #print('d_modname2id:',d_modname2id)
+            
         if seg_only:
-            images = {'Seg': results['G5']}
-            if 'G4' in results:
-                images.update({'Marker': results['G4']})
+            images = {'Seg': results[d_modname2id['Seg']]}
+            marker_key = find_marker_key(d_modname2id)
+            if marker_key is not None:
+                images[marker_key] = results[d_modname2id[marker_key]]
+            # if f'G{opt.modalities_no}' in results:
+            #     images.update({'Marker': results['G{opt.modalities_no}']})
+            # if 'Marker' in d_modname2id:
+            #     images.update({'Marker': results[d_modname2id['Marker']]})
         else:
-            images = {
-                'Hema': results['G1'],
-                'DAPI': results['G2'],
-                'Lap2': results['G3'],
-                'Marker': results['G4'],
-                'Seg': results['G5'],
-            }
+            # images = {
+            #     'Hema': results['G1'],
+            #     'DAPI': results['G2'],
+            #     'Lap2': results['G3'],
+            #     'Marker': results['G4'],
+            #     'Seg': results['G5'],
+            # }
+            # images = {f'mod{i+1}': results[f'G{i+1}'] for i in range(opt.modalities_no)}
+            # images['Seg'] = results[f'G{opt.modalities_no+1}']
+            images = {mod_name: results[mod_id] for mod_name,mod_id in d_modname2id.items()}
         
         if return_seg_intermediate and not seg_only:
-            images.update({'IHC_s':results['G51'],
-                          'Hema_s':results['G52'],
-                          'DAPI_s':results['G53'],
-                          'Lap2_s':results['G54'],
-                          'Marker_s':results['G55'],})
+            # images.update({'IHC_s':results['G51'],
+            #               'Hema_s':results['G52'],
+            #               'DAPI_s':results['G53'],
+            #               'Lap2_s':results['G54'],
+            #               'Marker_s':results['G55'],})
+            
+            #images.update({f'mod{i+1}_s':results[f'G{opt.modalities_no+1}{i+1}'] for i in range(opt.modalities_no+1)})
+            #images.update({f'{mod_name}_s':results[f'G{opt.modalities_no+1}']})
+            images.update({f'{mod_name}_s':results[d_modname2id_seg[mod_name]] for mod_name in d_modname2id_seg.keys()})
         
-        if color_dapi and not seg_only:
-            matrix = (       0,        0,        0, 0,
-                      299/1000, 587/1000, 114/1000, 0,
-                      299/1000, 587/1000, 114/1000, 0)
-            images['DAPI'] = images['DAPI'].convert('RGB', matrix)
-        if color_marker and not seg_only:
-            matrix = (299/1000, 587/1000, 114/1000, 0,
-                      299/1000, 587/1000, 114/1000, 0,
-                             0,        0,        0, 0)
-            images['Marker'] = images['Marker'].convert('RGB', matrix)
+        # if color_dapi and not seg_only:
+        #     matrix = (       0,        0,        0, 0,
+        #               299/1000, 587/1000, 114/1000, 0,
+        #               299/1000, 587/1000, 114/1000, 0)
+        #     images['DAPI'] = images['DAPI'].convert('RGB', matrix)
+        # if color_marker and not seg_only:
+        #     matrix = (299/1000, 587/1000, 114/1000, 0,
+        #               299/1000, 587/1000, 114/1000, 0,
+        #                      0,        0,        0, 0)
+        #     images['Marker'] = images['Marker'].convert('RGB', matrix)
         return images
 
     elif opt.model == 'DeepLIIFExt':
@@ -485,11 +564,11 @@ def inference(img, tile_size, overlap_size, model_path, use_torchserve=False,
         return results # return result images with default key names (i.e., net names)
 
 
-def postprocess(orig, images, tile_size, model, seg_thresh=150, size_thresh='default', marker_thresh=None, size_thresh_upper=None):
+def postprocess(orig, images, tile_size, model, seg_thresh=120, size_thresh='default', marker_thresh=None, size_thresh_upper=None):
     if model in ['DeepLIIF','DeepLIIFKD']:
         resolution = '40x' if tile_size > 384 else ('20x' if tile_size > 192 else '10x')
         overlay, refined, scoring = compute_final_results(
-            orig, images['Seg'], images.get('Marker'), resolution,
+            orig, images['Seg'], images.get(find_marker_key(images)), resolution,
             size_thresh, marker_thresh, size_thresh_upper, seg_thresh)
         processed_images = {}
         processed_images['SegOverlaid'] = Image.fromarray(overlay)
@@ -518,7 +597,7 @@ def postprocess(orig, images, tile_size, model, seg_thresh=150, size_thresh='def
 
 def infer_modalities(img, tile_size, model_dir, eager_mode=False,
                      color_dapi=False, color_marker=False, opt=None,
-                     return_seg_intermediate=False, seg_only=False, seg_weights=None):
+                     return_seg_intermediate=False, seg_only=False, mod_only=False, seg_weights=None):
     """
     This function is used to infer modalities for the given image using a trained model.
     :param img: The input image.
@@ -547,17 +626,21 @@ def infer_modalities(img, tile_size, model_dir, eager_mode=False,
         opt=opt,
         return_seg_intermediate=return_seg_intermediate,
         seg_only=seg_only,
+        mod_only=mod_only,
         seg_weights=seg_weights,
     )
 
     if not hasattr(opt,'seg_gen') or (hasattr(opt,'seg_gen') and opt.seg_gen): # the first condition accounts for old settings of deepliif; the second refers to deepliifext models
-        post_images, scoring = postprocess(img, images, tile_size, opt.model)
-        images = {**images, **post_images}
-        if seg_only:
-            delete_keys = [k for k in images.keys() if 'Seg' not in k]
-            for name in delete_keys:
-                del images[name]
-        return images, scoring
+        if not mod_only:
+            post_images, scoring = postprocess(img, images, tile_size, opt.model)
+            images = {**images, **post_images}
+            if seg_only:
+                delete_keys = [k for k in images.keys() if 'Seg' not in k]
+                for name in delete_keys:
+                    del images[name]
+            return images, scoring
+        else:
+            return images, None
     else:
         return images, None
 
@@ -763,7 +846,8 @@ def infer_cells_for_wsi(filename, model_dir, tile_size, region_size=20000, versi
 
                 seg = to_array(images['Seg'])
                 del images['Seg']
-                marker = to_array(images['Marker'], True) if 'Marker' in images else None
+                marker_key = find_marker_key(images)
+                marker = to_array(images[marker_key], True) if marker_key is not None else None
                 del images
                 region_data = compute_cell_results(seg, marker, resolution, version=version)
                 del seg
@@ -818,3 +902,10 @@ def infer_cells_for_wsi(filename, model_dir, tile_size, region_size=20000, versi
         data['modelVersion'] = 'unknown'
 
     return data
+
+
+def find_marker_key(dictionary):
+    for key in dictionary:
+        if key.endswith('Marker'):
+            return key
+    return None

deepliif/models/base_model.py

@@ -148,8 +148,13 @@ class BaseModel(ABC):
                 if not hasattr(self, name):
                     if len(name.split('_')) != 2:
                         if self.opt.model in ['DeepLIIF','DeepLIIFKD']:
-                            img_name = name[:-1] + '_' + name[-1]
-                            visual_ret[name] = getattr(self, img_name)
+                            if name.endswith('_teacher'):
+                                suffix = '_teacher'
+                                name = name[:-8]
+                            else:
+                                suffix = ''
+                            img_name = name[:-1] + '_' + name[-1] + suffix
+                            visual_ret[name+suffix] = getattr(self, img_name)
                         else:
                             if self.opt.model == 'CycleGAN':
                                 l_output = getattr(self, name.split('_')[0] + '_' + name.split('_')[1])

deepliif/options/__init__.py

@@ -2,7 +2,7 @@
 
 from pathlib import Path
 import os
-from ..util.util import mkdirs
+from ..util.util import mkdirs, init_input_and_mod_id
 import re
 
 def read_model_params(file_addr):
@@ -32,7 +32,7 @@ def read_model_params(file_addr):
             
             # if isinstance(param_dict[key],list):
             #     param_dict[key] = param_dict[key][0]
-            
+    
     return param_dict
 
 class Options:
@@ -43,7 +43,7 @@ class Options:
         
         if path_file:
             d_params = read_model_params(path_file)
-     
+        
         for k,v in d_params.items():
             try:
                 if k not in ['phase']: # e.g., k = 'phase', v = 'train', eval(v) is a function rather than a string
@@ -69,6 +69,7 @@ class Options:
         else:
             self.phase = 'test'
             self.is_train = False
+            self.continue_train = False
             self.input_nc = 3
             self.output_nc = 3
             self.ngf = 64
@@ -83,6 +84,41 @@ class Options:
                 # and can be configured in the inference function
                 self.BtoA = False if not hasattr(self,'BtoA') else self.BtoA 
             
+            # to account for old settings before modalities_no was introduced
+            if not hasattr(self,'modalities_no') and hasattr(self,'targets_no'):
+                self.modalities_no = self.targets_no - 1
+                del self.targets_no
+            
+            if self.model in ['DeepLIIF','DeepLIIFKD']:
+                self.mod_id_seg, self.input_id = init_input_and_mod_id(self, os.path.dirname(path_file))
+                self.input_id = int(self.input_id)
+                print('mod id seg:', self.mod_id_seg, '; input id:', self.input_id)
+            
+                print('Determining modalities names for test-mode model...')
+                if self.modalities_no == 4:
+                    if not hasattr(self,'modalities_names'):
+                        self.modalities_names = ['IHC','Hema','DAPI','Lap2','Marker']
+                        self.seg_weights = [0.5,0,0,0,0.5]
+                elif not hasattr(self,'modalities_names') or len(self.modalities_names)==0:
+                    # if self.model == 'DeepLIIFKD':
+                    #     # try find the modalities names from the teacher model
+                    #     d_params_teacher = read_model_params(os.path.join(self.model_dir_teacher,'train_opt.txt'))
+                    #     if 'modalities_names' in d_params_teacher:
+                    #         self.modalities_names = d_params_teacher['modalities_names']
+                    # # check again
+                    # if not hasattr(self,'modalities_names') or len(self.modalities_names)==0:
+                    self.modalities_names = [f'mod{i}' for i in range(self.modalities_no+1)]
+            else:
+                self.modalities_names = [f'mod{i}' for i in range(self.modalities_no+1)]
+
+            print('modalities names:', self.modalities_names)
+            
+            if not hasattr(self, 'background_colors'):
+                if self.model in ['DeepLIIF','DeepLIIFKD']:
+                    self.background_colors = [(201, 211, 208),(10, 10, 10), (0, 0, 0), (10, 10, 10)]
+                else:
+                    self.background_colors = [(10, 10, 10)] * self.modalities_no
+            
             # reset checkpoints_dir and name based on the model directory
             # when base model is initialized: self.save_dir = os.path.join(opt.checkpoints_dir, opt.name) 
             model_dir = Path(path_file).parent
@@ -95,15 +131,11 @@ class Options:
             if isinstance(self.gpu_ids,int):
                 self.gpu_ids = (self.gpu_ids,)
             
-            # to account for old settings before modalities_no was introduced
-            if not hasattr(self,'modalities_no') and hasattr(self,'targets_no'):
-                self.modalities_no = self.targets_no - 1
-                del self.targets_no
-            
             # to account for old settings: same as in cli.py train
             if not hasattr(self,'seg_no'):
               if self.model == 'DeepLIIF':
                   self.seg_no = 1
+                  self.seg_gen = True
               elif self.model == 'DeepLIIFExt':
                   if self.seg_gen:
                       self.seg_no = self.modalities_no

deepliif/postprocessing.py

@@ -80,7 +80,7 @@ def adjust_marker(inferred_tile, orig_tile):
 
 
 # Default postprocessing values
-DEFAULT_SEG_THRESH = 150
+DEFAULT_SEG_THRESH = 120
 DEFAULT_NOISE_THRESH = 4
 
 # Values for uint8 label masks