MPlexA 0.7.2__py3-none-any.whl

MPlexA/__init__.py

@@ -0,0 +1 @@
+__version__='0.7.2'

MPlexA/__main__.py

@@ -0,0 +1,12 @@
+from.import gui
+
+
+
+def main():
+
+	gui.main_window()
+
+
+if __name__=='__main__':
+
+	main()

MPlexA/detector.py

@@ -0,0 +1,169 @@
+import os
+import cv2
+import json
+import torch
+from MPlexA.detectron2 import model_zoo
+from MPlexA.detectron2.checkpoint import DetectionCheckpointer
+from MPlexA.detectron2.config import get_cfg
+from MPlexA.detectron2.data import MetadataCatalog,DatasetCatalog,build_detection_test_loader
+from MPlexA.detectron2.data.datasets import register_coco_instances
+from MPlexA.detectron2.engine import DefaultTrainer,DefaultPredictor
+from MPlexA.detectron2.utils.visualizer import Visualizer
+from MPlexA.detectron2.evaluation import COCOEvaluator,inference_on_dataset
+from MPlexA.detectron2.modeling import build_model
+
+
+
+class Detector():
+
+
+	def __init__(self):
+		self.device='cuda'if torch.cuda.is_available()else'cpu'# whether the GPU is available, if so, use GPU
+		self.cell_mapping=None# the celln categories and names in a Detector
+		self.inferencing_framesize=None
+		self.black_background=None
+		self.current_detector=None# the current Detector used for inference
+
+
+	def train(self,path_to_annotation,path_to_trainingimages,path_to_detector,iteration_num,inference_size,num_rois,black_background=0):
+		# path_to_annotation: the path to the .json file that stores the annotations in coco format
+		# path_to_trainingimages: the folder that stores all the training images
+		# iteration_num: the number of training iterations
+		# inference_size: the Detector inferencing frame size
+		# num_rois: the batch size of ROI heads per image
+		# black_background: whether the background of images to analyze is black/darker
+		if str('MPlexA_detector_train')in DatasetCatalog.list():
+			DatasetCatalog.remove('MPlexA_detector_train')
+			MetadataCatalog.remove('MPlexA_detector_train')
+		register_coco_instances('MPlexA_detector_train',{},path_to_annotation,path_to_trainingimages)
+		datasetcat=DatasetCatalog.get('MPlexA_detector_train')
+		metadatacat=MetadataCatalog.get('MPlexA_detector_train')
+		classnames=metadatacat.thing_classes
+		model_parameters_dict={}
+		model_parameters_dict['cell_names']=[]
+		annotation_data=json.load(open(path_to_annotation))
+		for i in annotation_data['categories']:
+			if i['id']>0:
+				model_parameters_dict['cell_names'].append(i['name'])
+		print('Cell names in annotation file: '+str(model_parameters_dict['cell_names']))
+		cfg=get_cfg()
+		cfg.merge_from_file(model_zoo.get_config_file('COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml'))
+		cfg.OUTPUT_DIR=path_to_detector
+		cfg.DATASETS.TRAIN=('MPlexA_detector_train',)
+		cfg.DATASETS.TEST=()
+		cfg.DATALOADER.NUM_WORKERS=4
+		cfg.MODEL.WEIGHTS=model_zoo.get_checkpoint_url('COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml')
+		cfg.MODEL.ROI_HEADS.BATCH_SIZE_PER_IMAGE=num_rois
+		cfg.MODEL.ROI_HEADS.NUM_CLASSES=int(len(classnames))
+		cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST=0.5
+		cfg.MODEL.DEVICE=self.device
+		cfg.SOLVER.IMS_PER_BATCH=4
+		cfg.SOLVER.MAX_ITER=int(iteration_num)
+		cfg.SOLVER.BASE_LR=0.001
+		cfg.SOLVER.WARMUP_ITERS=int(iteration_num*0.1)
+		cfg.SOLVER.STEPS=(int(iteration_num*0.4),int(iteration_num*0.8))
+		cfg.SOLVER.GAMMA=0.5
+		cfg.SOLVER.CHECKPOINT_PERIOD=100000000000000000
+		cfg.INPUT.MIN_SIZE_TEST=int(inference_size)
+		cfg.INPUT.MAX_SIZE_TEST=int(inference_size)
+		cfg.INPUT.MIN_SIZE_TRAIN=(int(inference_size),)
+		cfg.INPUT.MAX_SIZE_TRAIN=int(inference_size)
+		os.makedirs(cfg.OUTPUT_DIR)
+		trainer=DefaultTrainer(cfg)
+		trainer.resume_or_load(False)
+		trainer.train()
+		model_parameters=os.path.join(cfg.OUTPUT_DIR,'model_parameters.txt')
+		model_parameters_dict['cell_mapping']={}
+		model_parameters_dict['inferencing_framesize']=int(inference_size)
+		model_parameters_dict['black_background']=int(black_background)
+		for i in range(len(classnames)):
+			model_parameters_dict['cell_mapping'][i]=classnames[i]
+		with open(model_parameters,'w')as f:
+			f.write(json.dumps(model_parameters_dict))
+		predictor=DefaultPredictor(cfg)
+		model=predictor.model
+		DetectionCheckpointer(model).resume_or_load(os.path.join(cfg.OUTPUT_DIR,'model_final.pth'))
+		model.eval()
+		config=os.path.join(cfg.OUTPUT_DIR,'config.yaml')
+		with open(config,'w')as f:
+			f.write(cfg.dump())
+		print('Detector training completed!')
+
+
+	def test(self,path_to_annotation,path_to_testingimages,path_to_detector,output_path):
+		# path_to_annotation: the path to the .json file that stores the annotations in coco format
+		# path_to_testingimages: the folder that stores all the ground-truth testing images
+		# output_path: the folder that stores the testing images with annotations
+		if str('MPlexA_detector_test')in DatasetCatalog.list():
+			DatasetCatalog.remove('MPlexA_detector_test')
+			MetadataCatalog.remove('MPlexA_detector_test')
+		register_coco_instances('MPlexA_detector_test',{},path_to_annotation,path_to_testingimages)
+		datasetcat=DatasetCatalog.get('MPlexA_detector_test')
+		metadatacat=MetadataCatalog.get('MPlexA_detector_test')
+		cellmapping=os.path.join(path_to_detector,'model_parameters.txt')
+		with open(cellmapping)as f:
+			model_parameters=f.read()
+		cell_names=json.loads(model_parameters)['cell_names']
+		dt_infersize=int(json.loads(model_parameters)['inferencing_framesize'])
+		bg=int(json.loads(model_parameters)['black_background'])
+		print('The total categories of cells in this Detector: '+str(cell_names))
+		print('The inferencing framesize of this Detector: '+str(dt_infersize))
+		if bg==0:
+			print('The images that can be analyzed by this Detector have black/darker background')
+		else:
+			print('The images that can be analyzed by this Detector have white/lighter background')
+		cfg=get_cfg()
+		cfg.set_new_allowed(True)
+		cfg.merge_from_file(os.path.join(path_to_detector,'config.yaml'))
+		cfg.MODEL.WEIGHTS=os.path.join(path_to_detector,'model_final.pth')
+		cfg.MODEL.DEVICE=self.device
+		predictor=DefaultPredictor(cfg)
+		for d in datasetcat:
+			im=cv2.imread(d['file_name'])
+			outputs=predictor(im)
+			v=Visualizer(im[:,:,::-1],MetadataCatalog.get('MPlexA_detector_test'),scale=1.2)
+			out=v.draw_instance_predictions(outputs['instances'].to('cpu'))
+			cv2.imwrite(os.path.join(output_path,os.path.basename(d['file_name'])),out.get_image()[:,:,::-1])
+		evaluator=COCOEvaluator('MPlexA_detector_test',cfg,False,output_dir=output_path)
+		val_loader=build_detection_test_loader(cfg,'MPlexA_detector_test')
+		inference_on_dataset(predictor.model,val_loader,evaluator)
+		mAP=evaluator._results['bbox']['AP']
+		print(f'The mean average precision (mAP) of the Detector is: {mAP:.4f}%.')
+		print('Detector testing completed!')
+
+
+	def load(self,path_to_detector,cell_kinds):
+		# cell_kinds: the catgories of cells / objects to be analyzed
+		config=os.path.join(path_to_detector,'config.yaml')
+		detector_model=os.path.join(path_to_detector,'model_final.pth')
+		cellmapping=os.path.join(path_to_detector,'model_parameters.txt')
+		with open(cellmapping)as f:
+			model_parameters=f.read()
+		self.cell_mapping=json.loads(model_parameters)['cell_mapping']
+		cell_names=json.loads(model_parameters)['cell_names']
+		self.inferencing_framesize=int(json.loads(model_parameters)['inferencing_framesize'])
+		bg=int(json.loads(model_parameters)['black_background'])
+		print('The total categories of cells in this Detector: '+str(cell_names))
+		print('The cells of interest in this Detector: '+str(cell_kinds))
+		print('The inferencing framesize of this Detector: '+str(self.inferencing_framesize))
+		if bg==0:
+			self.black_background=True
+			print('The images that can be analyzed by this Detector have black/darker background')
+		else:
+			self.black_background=False
+			print('The images that can be analyzed by this Detector have white/lighter background')
+		cfg=get_cfg()
+		cfg.set_new_allowed(True)
+		cfg.merge_from_file(config)
+		cfg.MODEL.DEVICE=self.device
+		self.current_detector=build_model(cfg)
+		DetectionCheckpointer(self.current_detector).load(detector_model)
+		self.current_detector.eval()
+
+
+	def inference(self,inputs):
+		# inputs: images that the current Detector runs on
+		with torch.no_grad():
+			outputs=self.current_detector(inputs)
+		return outputs
+

MPlexA/detectron2/__init__.py

@@ -0,0 +1,6 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+from.utils.env import setup_environment
+setup_environment()
+# This line will be programatically read/write by setup.py.
+# Leave them at the bottom of this file and don't touch them.
+__version__='0.6'

MPlexA/detectron2/checkpoint/__init__.py

@@ -0,0 +1,7 @@
+# -*- coding: utf-8 -*-
+# Copyright (c) Facebook, Inc. and its affiliates.
+# File:
+from.import catalog as _UNUSED# register the handler
+from.detection_checkpoint import DetectionCheckpointer
+from fvcore.common.checkpoint import Checkpointer,PeriodicCheckpointer
+__all__=['Checkpointer','PeriodicCheckpointer','DetectionCheckpointer']

MPlexA/detectron2/checkpoint/c2_model_loading.py

@@ -0,0 +1,382 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+import copy
+import logging
+import re
+from typing import Dict,List
+import torch
+
+
+def convert_basic_c2_names(original_keys):
+	'''
+    Apply some basic name conversion to names in C2 weights.
+    It only deals with typical backbone models.
+
+    Args:
+        original_keys (list[str]):
+    Returns:
+        list[str]: The same number of strings matching those in original_keys.
+    '''
+	layer_keys=copy.deepcopy(original_keys)
+	layer_keys=[
+		{'pred_b':'linear_b','pred_w':'linear_w'}.get(k,k)for k in layer_keys
+	]# some hard-coded mappings
+	layer_keys=[k.replace('_','.')for k in layer_keys]
+	layer_keys=[re.sub('\\.b$','.bias',k)for k in layer_keys]
+	layer_keys=[re.sub('\\.w$','.weight',k)for k in layer_keys]
+	# Uniform both bn and gn names to "norm"
+	layer_keys=[re.sub('bn\\.s$','norm.weight',k)for k in layer_keys]
+	layer_keys=[re.sub('bn\\.bias$','norm.bias',k)for k in layer_keys]
+	layer_keys=[re.sub('bn\\.rm','norm.running_mean',k)for k in layer_keys]
+	layer_keys=[re.sub('bn\\.running.mean$','norm.running_mean',k)for k in layer_keys]
+	layer_keys=[re.sub('bn\\.riv$','norm.running_var',k)for k in layer_keys]
+	layer_keys=[re.sub('bn\\.running.var$','norm.running_var',k)for k in layer_keys]
+	layer_keys=[re.sub('bn\\.gamma$','norm.weight',k)for k in layer_keys]
+	layer_keys=[re.sub('bn\\.beta$','norm.bias',k)for k in layer_keys]
+	layer_keys=[re.sub('gn\\.s$','norm.weight',k)for k in layer_keys]
+	layer_keys=[re.sub('gn\\.bias$','norm.bias',k)for k in layer_keys]
+	# stem
+	layer_keys=[re.sub('^res\\.conv1\\.norm\\.','conv1.norm.',k)for k in layer_keys]
+	# to avoid mis-matching with "conv1" in other components (e.g. detection head)
+	layer_keys=[re.sub('^conv1\\.','stem.conv1.',k)for k in layer_keys]
+	# layer1-4 is used by torchvision, however we follow the C2 naming strategy (res2-5)
+	# layer_keys = [re.sub("^res2.", "layer1.", k) for k in layer_keys]
+	# layer_keys = [re.sub("^res3.", "layer2.", k) for k in layer_keys]
+	# layer_keys = [re.sub("^res4.", "layer3.", k) for k in layer_keys]
+	# layer_keys = [re.sub("^res5.", "layer4.", k) for k in layer_keys]
+	# blocks
+	layer_keys=[k.replace('.branch1.','.shortcut.')for k in layer_keys]
+	layer_keys=[k.replace('.branch2a.','.conv1.')for k in layer_keys]
+	layer_keys=[k.replace('.branch2b.','.conv2.')for k in layer_keys]
+	layer_keys=[k.replace('.branch2c.','.conv3.')for k in layer_keys]
+	# DensePose substitutions
+	layer_keys=[re.sub('^body.conv.fcn','body_conv_fcn',k)for k in layer_keys]
+	layer_keys=[k.replace('AnnIndex.lowres','ann_index_lowres')for k in layer_keys]
+	layer_keys=[k.replace('Index.UV.lowres','index_uv_lowres')for k in layer_keys]
+	layer_keys=[k.replace('U.lowres','u_lowres')for k in layer_keys]
+	layer_keys=[k.replace('V.lowres','v_lowres')for k in layer_keys]
+	return layer_keys
+
+
+def convert_c2_detectron_names(weights):
+	'''
+    Map Caffe2 Detectron weight names to Detectron2 names.
+
+    Args:
+        weights (dict): name -> tensor
+
+    Returns:
+        dict: detectron2 names -> tensor
+        dict: detectron2 names -> C2 names
+    '''
+	logger=logging.getLogger(__name__)
+	logger.info('Renaming Caffe2 weights ......')
+	original_keys=sorted(weights.keys())
+	layer_keys=copy.deepcopy(original_keys)
+	layer_keys=convert_basic_c2_names(layer_keys)
+	# --------------------------------------------------------------------------
+	# RPN hidden representation conv
+	# --------------------------------------------------------------------------
+	# FPN case
+	# In the C2 model, the RPN hidden layer conv is defined for FPN level 2 and then
+	# shared for all other levels, hence the appearance of "fpn2"
+	layer_keys=[
+		k.replace('conv.rpn.fpn2','proposal_generator.rpn_head.conv')for k in layer_keys
+	]
+	# Non-FPN case
+	layer_keys=[k.replace('conv.rpn','proposal_generator.rpn_head.conv')for k in layer_keys]
+	# --------------------------------------------------------------------------
+	# RPN box transformation conv
+	# --------------------------------------------------------------------------
+	# FPN case (see note above about "fpn2")
+	layer_keys=[
+		k.replace('rpn.bbox.pred.fpn2','proposal_generator.rpn_head.anchor_deltas')
+		for k in layer_keys
+	]
+	layer_keys=[
+		k.replace('rpn.cls.logits.fpn2','proposal_generator.rpn_head.objectness_logits')
+		for k in layer_keys
+	]
+	# Non-FPN case
+	layer_keys=[
+		k.replace('rpn.bbox.pred','proposal_generator.rpn_head.anchor_deltas')for k in layer_keys
+	]
+	layer_keys=[
+		k.replace('rpn.cls.logits','proposal_generator.rpn_head.objectness_logits')
+		for k in layer_keys
+	]
+	# --------------------------------------------------------------------------
+	# Fast R-CNN box head
+	# --------------------------------------------------------------------------
+	layer_keys=[re.sub('^bbox\\.pred','bbox_pred',k)for k in layer_keys]
+	layer_keys=[re.sub('^cls\\.score','cls_score',k)for k in layer_keys]
+	layer_keys=[re.sub('^fc6\\.','box_head.fc1.',k)for k in layer_keys]
+	layer_keys=[re.sub('^fc7\\.','box_head.fc2.',k)for k in layer_keys]
+	# 4conv1fc head tensor names: head_conv1_w, head_conv1_gn_s
+	layer_keys=[re.sub('^head\\.conv','box_head.conv',k)for k in layer_keys]
+	# --------------------------------------------------------------------------
+	# FPN lateral and output convolutions
+	# --------------------------------------------------------------------------
+
+
+	def fpn_map(name):
+		'''
+        Look for keys with the following patterns:
+        1) Starts with "fpn.inner."
+           Example: "fpn.inner.res2.2.sum.lateral.weight"
+           Meaning: These are lateral pathway convolutions
+        2) Starts with "fpn.res"
+           Example: "fpn.res2.2.sum.weight"
+           Meaning: These are FPN output convolutions
+        '''
+		splits=name.split('.')
+		norm='.norm'if'norm'in splits else''
+		if name.startswith('fpn.inner.'):
+			# splits example: ['fpn', 'inner', 'res2', '2', 'sum', 'lateral', 'weight']
+			stage=int(splits[2][len('res'):])
+			return'fpn_lateral{}{}.{}'.format(stage,norm,splits[-1])
+		elif name.startswith('fpn.res'):
+			# splits example: ['fpn', 'res2', '2', 'sum', 'weight']
+			stage=int(splits[1][len('res'):])
+			return'fpn_output{}{}.{}'.format(stage,norm,splits[-1])
+		return name
+	layer_keys=[fpn_map(k)for k in layer_keys]
+	# --------------------------------------------------------------------------
+	# Mask R-CNN mask head
+	# --------------------------------------------------------------------------
+	# roi_heads.StandardROIHeads case
+	layer_keys=[k.replace('.[mask].fcn','mask_head.mask_fcn')for k in layer_keys]
+	layer_keys=[re.sub('^\\.mask\\.fcn','mask_head.mask_fcn',k)for k in layer_keys]
+	layer_keys=[k.replace('mask.fcn.logits','mask_head.predictor')for k in layer_keys]
+	# roi_heads.Res5ROIHeads case
+	layer_keys=[k.replace('conv5.mask','mask_head.deconv')for k in layer_keys]
+	# --------------------------------------------------------------------------
+	# Keypoint R-CNN head
+	# --------------------------------------------------------------------------
+	# interestingly, the keypoint head convs have blob names that are simply "conv_fcnX"
+	layer_keys=[k.replace('conv.fcn','roi_heads.keypoint_head.conv_fcn')for k in layer_keys]
+	layer_keys=[
+		k.replace('kps.score.lowres','roi_heads.keypoint_head.score_lowres')for k in layer_keys
+	]
+	layer_keys=[k.replace('kps.score.','roi_heads.keypoint_head.score.')for k in layer_keys]
+	# --------------------------------------------------------------------------
+	# Done with replacements
+	# --------------------------------------------------------------------------
+	assert len(set(layer_keys))==len(layer_keys)
+	assert len(original_keys)==len(layer_keys)
+	new_weights={}
+	new_keys_to_original_keys={}
+	for orig,renamed in zip(original_keys,layer_keys):
+		new_keys_to_original_keys[renamed]=orig
+		if renamed.startswith('bbox_pred.')or renamed.startswith('mask_head.predictor.'):
+			# remove the meaningless prediction weight for background class
+			new_start_idx=4 if renamed.startswith('bbox_pred.')else 1
+			new_weights[renamed]=weights[orig][new_start_idx:]
+			logger.info(
+				'Remove prediction weight for background class in {}. The shape changes from '
+				'{} to {}.'.format(
+					renamed,tuple(weights[orig].shape),tuple(new_weights[renamed].shape)
+				)
+			)
+		elif renamed.startswith('cls_score.'):
+			# move weights of bg class from original index 0 to last index
+			logger.info(
+				'Move classification weights for background class in {} from index 0 to '
+				'index {}.'.format(renamed,weights[orig].shape[0]-1)
+			)
+			new_weights[renamed]=torch.cat([weights[orig][1:],weights[orig][:1]])
+		else:
+			new_weights[renamed]=weights[orig]
+	return new_weights,new_keys_to_original_keys
+# Note the current matching is not symmetric.
+# it assumes model_state_dict will have longer names.
+
+
+def align_and_update_state_dicts(model_state_dict,ckpt_state_dict,c2_conversion=True):
+	'''
+    Match names between the two state-dict, and returns a new chkpt_state_dict with names
+    converted to match model_state_dict with heuristics. The returned dict can be later
+    loaded with fvcore checkpointer.
+    If `c2_conversion==True`, `ckpt_state_dict` is assumed to be a Caffe2
+    model and will be renamed at first.
+
+    Strategy: suppose that the models that we will create will have prefixes appended
+    to each of its keys, for example due to an extra level of nesting that the original
+    pre-trained weights from ImageNet won't contain. For example, model.state_dict()
+    might return backbone[0].body.res2.conv1.weight, while the pre-trained model contains
+    res2.conv1.weight. We thus want to match both parameters together.
+    For that, we look for each model weight, look among all loaded keys if there is one
+    that is a suffix of the current weight name, and use it if that's the case.
+    If multiple matches exist, take the one with longest size
+    of the corresponding name. For example, for the same model as before, the pretrained
+    weight file can contain both res2.conv1.weight, as well as conv1.weight. In this case,
+    we want to match backbone[0].body.conv1.weight to conv1.weight, and
+    backbone[0].body.res2.conv1.weight to res2.conv1.weight.
+    '''
+	model_keys=sorted(model_state_dict.keys())
+	if c2_conversion:
+		ckpt_state_dict,original_keys=convert_c2_detectron_names(ckpt_state_dict)
+		# original_keys: the name in the original dict (before renaming)
+	else:
+		original_keys={x:x for x in ckpt_state_dict.keys()}
+	ckpt_keys=sorted(ckpt_state_dict.keys())
+
+
+	def match(a,b):
+		# Matched ckpt_key should be a complete (starts with '.') suffix.
+		# For example, roi_heads.mesh_head.whatever_conv1 does not match conv1,
+		# but matches whatever_conv1 or mesh_head.whatever_conv1.
+		return a==b or a.endswith('.'+b)
+	# get a matrix of string matches, where each (i, j) entry correspond to the size of the
+	# ckpt_key string, if it matches
+	match_matrix=[len(j)if match(i,j)else 0 for i in model_keys for j in ckpt_keys]
+	match_matrix=torch.as_tensor(match_matrix).view(len(model_keys),len(ckpt_keys))
+	# use the matched one with longest size in case of multiple matches
+	max_match_size,idxs=match_matrix.max(1)
+	# remove indices that correspond to no-match
+	idxs[max_match_size==0]=-1
+	logger=logging.getLogger(__name__)
+	# matched_pairs (matched checkpoint key --> matched model key)
+	matched_keys={}
+	result_state_dict={}
+	for idx_model,idx_ckpt in enumerate(idxs.tolist()):
+		if idx_ckpt==-1:
+			continue
+		key_model=model_keys[idx_model]
+		key_ckpt=ckpt_keys[idx_ckpt]
+		value_ckpt=ckpt_state_dict[key_ckpt]
+		shape_in_model=model_state_dict[key_model].shape
+		if shape_in_model!=value_ckpt.shape:
+			logger.warning(
+				'Shape of {} in checkpoint is {}, while shape of {} in model is {}.'.format(
+					key_ckpt,value_ckpt.shape,key_model,shape_in_model
+				)
+			)
+			logger.warning(
+				'{} will not be loaded. Please double check and see if this is desired.'.format(
+					key_ckpt
+				)
+			)
+			continue
+		assert key_model not in result_state_dict
+		result_state_dict[key_model]=value_ckpt
+		if key_ckpt in matched_keys:# already added to matched_keys
+			logger.error(
+				'Ambiguity found for {} in checkpoint!'
+				'It matches at least two keys in the model ({} and {}).'.format(
+					key_ckpt,key_model,matched_keys[key_ckpt]
+				)
+			)
+			raise ValueError('Cannot match one checkpoint key to multiple keys in the model.')
+		matched_keys[key_ckpt]=key_model
+	# logging:
+	matched_model_keys=sorted(matched_keys.values())
+	if len(matched_model_keys)==0:
+		logger.warning('No weights in checkpoint matched with model.')
+		return ckpt_state_dict
+	common_prefix=_longest_common_prefix(matched_model_keys)
+	rev_matched_keys={v:k for k,v in matched_keys.items()}
+	original_keys={k:original_keys[rev_matched_keys[k]]for k in matched_model_keys}
+	model_key_groups=_group_keys_by_module(matched_model_keys,original_keys)
+	table=[]
+	memo=set()
+	for key_model in matched_model_keys:
+		if key_model in memo:
+			continue
+		if key_model in model_key_groups:
+			group=model_key_groups[key_model]
+			memo|=set(group)
+			shapes=[tuple(model_state_dict[k].shape)for k in group]
+			table.append(
+				(
+					_longest_common_prefix([k[len(common_prefix):]for k in group])+'*',
+					_group_str([original_keys[k]for k in group]),
+					' '.join([str(x).replace(' ','')for x in shapes]),
+				)
+			)
+		else:
+			key_checkpoint=original_keys[key_model]
+			shape=str(tuple(model_state_dict[key_model].shape))
+			table.append((key_model[len(common_prefix):],key_checkpoint,shape))
+	submodule_str=common_prefix[:-1]if common_prefix else'model'
+	logger.info(
+		f'Following weights matched with submodule {submodule_str} - Total num: {len(table)}'
+	)
+	unmatched_ckpt_keys=[k for k in ckpt_keys if k not in set(matched_keys.keys())]
+	for k in unmatched_ckpt_keys:
+		result_state_dict[k]=ckpt_state_dict[k]
+	return result_state_dict
+
+
+def _group_keys_by_module(keys:List[str],original_names:Dict[str,str]):
+	'''
+    Params in the same submodule are grouped together.
+
+    Args:
+        keys: names of all parameters
+        original_names: mapping from parameter name to their name in the checkpoint
+
+    Returns:
+        dict[name -> all other names in the same group]
+    '''
+
+
+	def _submodule_name(key):
+		pos=key.rfind('.')
+		if pos<0:
+			return None
+		prefix=key[:pos+1]
+		return prefix
+	all_submodules=[_submodule_name(k)for k in keys]
+	all_submodules=[x for x in all_submodules if x]
+	all_submodules=sorted(all_submodules,key=len)
+	ret={}
+	for prefix in all_submodules:
+		group=[k for k in keys if k.startswith(prefix)]
+		if len(group)<=1:
+			continue
+		original_name_lcp=_longest_common_prefix_str([original_names[k]for k in group])
+		if len(original_name_lcp)==0:
+			# don't group weights if original names don't share prefix
+			continue
+		for k in group:
+			if k in ret:
+				continue
+			ret[k]=group
+	return ret
+
+
+def _longest_common_prefix(names:List[str])->str:
+	'''
+    ["abc.zfg", "abc.zef"] -> "abc."
+    '''
+	names=[n.split('.')for n in names]
+	m1,m2=min(names),max(names)
+	ret=[a for a,b in zip(m1,m2)if a==b]
+	ret='.'.join(ret)+'.'if len(ret)else''
+	return ret
+
+
+def _longest_common_prefix_str(names:List[str])->str:
+	m1,m2=min(names),max(names)
+	lcp=[]
+	for a,b in zip(m1,m2):
+		if a==b:
+			lcp.append(a)
+		else:
+			break
+	lcp=''.join(lcp)
+	return lcp
+
+
+def _group_str(names:List[str])->str:
+	'''
+    Turn "common1", "common2", "common3" into "common{1,2,3}"
+    '''
+	lcp=_longest_common_prefix_str(names)
+	rest=[x[len(lcp):]for x in names]
+	rest='{'+','.join(rest)+'}'
+	ret=lcp+rest
+	# add some simplification for BN specifically
+	ret=ret.replace('bn_{beta,running_mean,running_var,gamma}','bn_*')
+	ret=ret.replace('bn_beta,bn_running_mean,bn_running_var,bn_gamma','bn_*')
+	return ret