support showing model statistics in readthedocs (#262)

* support showing model statistics in readthedocs

* fix a bug

* fix a bug

* chmod 775 stat.py

docs/conf.py

@@ -12,6 +12,7 @@
 # documentation root, use os.path.abspath to make it absolute, like shown here.
 #
 import os
+import subprocess
 import sys
 
 sys.path.insert(0, os.path.abspath('..'))
@@ -84,8 +85,10 @@ def get_version():
 
 language = 'en'
 
-# def builder_inited_handler(app):
-#     subprocess.run(['./stat.py'])
 
-# def setup(app):
-#     app.connect('builder-inited', builder_inited_handler)
+def builder_inited_handler(app):
+    subprocess.run(['./stat.py'])
+
+
+def setup(app):
+    app.connect('builder-inited', builder_inited_handler)

docs/index.rst

@@ -8,6 +8,7 @@ You can switch between Chinese and English documents in the lower-left corner of
    :caption: Get Started
 
    install.md
+   modelzoo_statistics.md
    model_zoo.md
 
 .. toctree::

docs/stat.py

@@ -0,0 +1,62 @@
+#!/usr/bin/env python
+# Copyright (c) OpenMMLab. All rights reserved.
+import functools as func
+import glob
+import os.path as osp
+import re
+
+import numpy as np
+
+url_prefix = 'https://github.com/open-mmlab/mmtracking/blob/master/'
+
+files = sorted(glob.glob('../configs/*/*/README.md'))
+files.extend(glob.glob('../configs/*/README.md'))
+
+stats = []
+titles = []
+num_ckpts = 0
+
+for f in files:
+    url = osp.dirname(f.replace('../', url_prefix))
+
+    with open(f, 'r') as content_file:
+        content = content_file.read()
+
+    title = content.split('\n')[0].replace('#', '').strip()
+    ckpts = set(x.lower().strip()
+                for x in re.findall(r'https?://download.*\.pth', content)
+                if 'mmtracking' in x)
+    if len(ckpts) == 0:
+        continue
+
+    _papertype = [x for x in re.findall(r'\s*\[([A-Z]*?)\]\s*', content)]
+    assert len(_papertype) > 0
+    papertype = _papertype[0]
+
+    paper = set([(papertype, title)])
+
+    titles.append(title)
+    num_ckpts += len(ckpts)
+    statsmsg = f"""
+\t* [{papertype}] [{title}]({url}) ({len(ckpts)} ckpts)
+"""
+    stats.append((paper, ckpts, statsmsg))
+
+allpapers = func.reduce(lambda a, b: a.union(b), [p for p, _, _ in stats])
+msglist = '\n'.join(x for _, _, x in stats)
+
+papertypes, papercounts = np.unique([t for t, _ in allpapers],
+                                    return_counts=True)
+countstr = '\n'.join(
+    [f'   - {t}: {c}' for t, c in zip(papertypes, papercounts)])
+
+modelzoo = f"""
+# Model Zoo Statistics
+* Number of papers: {len(set(titles))}
+{countstr}
+* Number of checkpoints: {num_ckpts}
+{msglist}
+"""
+
+with open('modelzoo_statistics.md', 'w') as f:
+    f.write(modelzoo)

docs_zh-CN/conf.py

@@ -12,6 +12,7 @@
 # documentation root, use os.path.abspath to make it absolute, like shown here.
 #
 import os
+import subprocess
 import sys
 
 sys.path.insert(0, os.path.abspath('..'))
@@ -84,8 +85,10 @@ def get_version():
 
 language = 'zh_CN'
 
-# def builder_inited_handler(app):
-#     subprocess.run(['./stat.py'])
 
-# def setup(app):
-#     app.connect('builder-inited', builder_inited_handler)
+def builder_inited_handler(app):
+    subprocess.run(['./stat.py'])
+
+
+def setup(app):
+    app.connect('builder-inited', builder_inited_handler)

docs_zh-CN/index.rst

@@ -8,6 +8,7 @@
    :caption: 开始你的第一步
 
    install.md
+   modelzoo_statistics.md
    model_zoo.md
 
 .. toctree::

docs_zh-CN/stat.py

@@ -0,0 +1,62 @@
+#!/usr/bin/env python
+# Copyright (c) OpenMMLab. All rights reserved.
+import functools as func
+import glob
+import os.path as osp
+import re
+
+import numpy as np
+
+url_prefix = 'https://github.com/open-mmlab/mmtracking/blob/master/'
+
+files = sorted(glob.glob('../configs/*/*/README.md'))
+files.extend(glob.glob('../configs/*/README.md'))
+
+stats = []
+titles = []
+num_ckpts = 0
+
+for f in files:
+    url = osp.dirname(f.replace('../', url_prefix))
+
+    with open(f, 'r') as content_file:
+        content = content_file.read()
+
+    title = content.split('\n')[0].replace('#', '').strip()
+    ckpts = set(x.lower().strip()
+                for x in re.findall(r'https?://download.*\.pth', content)
+                if 'mmtracking' in x)
+    if len(ckpts) == 0:
+        continue
+
+    _papertype = [x for x in re.findall(r'\s*\[([A-Z]*?)\]\s*', content)]
+    assert len(_papertype) > 0
+    papertype = _papertype[0]
+
+    paper = set([(papertype, title)])
+
+    titles.append(title)
+    num_ckpts += len(ckpts)
+    statsmsg = f"""
+\t* [{papertype}] [{title}]({url}) ({len(ckpts)} ckpts)
+"""
+    stats.append((paper, ckpts, statsmsg))
+
+allpapers = func.reduce(lambda a, b: a.union(b), [p for p, _, _ in stats])
+msglist = '\n'.join(x for _, _, x in stats)
+
+papertypes, papercounts = np.unique([t for t, _ in allpapers],
+                                    return_counts=True)
+countstr = '\n'.join(
+    [f'   - {t}: {c}' for t, c in zip(papertypes, papercounts)])
+
+modelzoo = f"""
+# 模型库统计数据
+* 论文数量: {len(set(titles))}
+{countstr}
+* 模型数量: {num_ckpts}
+{msglist}
+"""
+
+with open('modelzoo_statistics.md', 'w') as f:
+    f.write(modelzoo)