1212import re
1313import sys
1414
15- import mmcv
1615from lxml import etree
16+ from mmcv .fileio import dump
1717
1818MMSEG_ROOT = osp .dirname (osp .dirname ((osp .dirname (__file__ ))))
1919
20+ COLLECTIONS = [
21+ 'ANN' , 'APCNet' , 'BiSeNetV1' , 'BiSeNetV2' , 'CCNet' , 'CGNet' , 'DANet' ,
22+ 'DeepLabV3' , 'DeepLabV3+' , 'DMNet' , 'DNLNet' , 'DPT' , 'EMANet' , 'EncNet' ,
23+ 'ERFNet' , 'FastFCN' , 'FastSCNN' , 'FCN' , 'GCNet' , 'ICNet' , 'ISANet' , 'KNet' ,
24+ 'NonLocalNet' , 'OCRNet' , 'PointRend' , 'PSANet' , 'PSPNet' , 'Segformer' ,
25+ 'Segmenter' , 'FPN' , 'SETR' , 'STDC' , 'UNet' , 'UPerNet'
26+ ]
27+ COLLECTIONS_TEMP = []
28+
2029
2130def dump_yaml_and_check_difference (obj , filename , sort_keys = False ):
2231 """Dump object to a yaml file, and check if the file content is different
@@ -30,7 +39,7 @@ def dump_yaml_and_check_difference(obj, filename, sort_keys=False):
3039 Bool: If the target YAML file is different from the original.
3140 """
3241
33- str_dump = mmcv . dump (obj , None , file_format = 'yaml' , sort_keys = sort_keys )
42+ str_dump = dump (obj , None , file_format = 'yaml' , sort_keys = sort_keys )
3443 if osp .isfile (filename ):
3544 file_exists = True
3645 with open (filename , 'r' , encoding = 'utf-8' ) as f :
@@ -88,7 +97,7 @@ def parse_md(md_file):
8897 # should be set with head or neck of this config file.
8998 is_backbone = None
9099
91- with open (md_file , 'r' ) as md :
100+ with open (md_file , 'r' , encoding = 'UTF-8' ) as md :
92101 lines = md .readlines ()
93102 i = 0
94103 current_dataset = ''
@@ -127,8 +136,9 @@ def parse_md(md_file):
127136 elif line [:15 ] == '<!-- [BACKBONE]' :
128137 is_backbone = True
129138 i += 1
130- elif line [0 ] == '|' and (
131- i + 1 ) < len (lines ) and lines [i + 1 ][:3 ] == '| -' :
139+ elif (line [0 ] == '|' and (i + 1 ) < len (lines )
140+ and lines [i + 1 ][:3 ] == '| -' and 'Method' in line
141+ and 'Crop Size' in line and 'Mem (GB)' in line ):
132142 cols = [col .strip () for col in line .split ('|' )]
133143 method_id = cols .index ('Method' )
134144 backbone_id = cols .index ('Backbone' )
@@ -246,11 +256,21 @@ def parse_md(md_file):
246256 collection .pop (check_key )
247257 else :
248258 collection [check_key ].pop (key )
259+ yml_file = f'{ md_file [:- 9 ]} { collection_name }
249260 if is_backbone :
250- result = {'Models' : models }
261+ if collection ['Name' ] not in COLLECTIONS :
262+ result = {
263+ 'Collections' : [collection ],
264+ 'Models' : models ,
265+ 'Yml' : yml_file
266+ }
267+ COLLECTIONS_TEMP .append (result )
268+ return False
269+ else :
270+ result = {'Models' : models }
251271 else :
272+ COLLECTIONS .append (collection ['Name' ])
252273 result = {'Collections' : [collection ], 'Models' : models }
253- yml_file = f'{ md_file [:- 9 ]} { collection_name }
254274 return dump_yaml_and_check_difference (result , yml_file )
255275
256276
@@ -286,6 +306,12 @@ def update_model_index():
286306 for fn in file_list :
287307 file_modified |= parse_md (fn )
288308
289- file_modified |= update_model_index ()
309+ for result in COLLECTIONS_TEMP :
310+ collection = result ['Collections' ][0 ]
311+ yml_file = result .pop ('Yml' , None )
312+ if collection ['Name' ] in COLLECTIONS :
313+ result .pop ('Collections' )
314+ file_modified |= dump_yaml_and_check_difference (result , yml_file )
290315
316+ file_modified |= update_model_index ()
291317 sys .exit (1 if file_modified else 0 )
0 commit comments