Spaces:
Running
Running
| import os | |
| import re | |
| import json | |
| import requests | |
| import gradio as gr | |
| import pandas as pd | |
| from tqdm import tqdm | |
| from bs4 import BeautifulSoup | |
| cache_json = "cv_backbones.json" | |
| def parse_url(url): | |
| response = requests.get(url) | |
| html = response.text | |
| return BeautifulSoup(html, "html.parser") | |
| def special_type(m_ver): | |
| m_type = re.search("[a-zA-Z]+", m_ver).group(0) | |
| if m_type == "wide" or m_type == "resnext": | |
| return "resnet" | |
| elif m_type == "swin": | |
| return "swin_transformer" | |
| elif m_type == "inception": | |
| return "googlenet" | |
| return m_type | |
| def info_on_dataset(m_ver, m_type, in1k_span): | |
| url_span = in1k_span.find_next_sibling("span", {"class": "s2"}) | |
| size_span = url_span.find_next_sibling("span", {"class": "mi"}) | |
| m_url = str(url_span.text[1:-1]) | |
| input_size = int(size_span.text) | |
| m_dict = {"ver": m_ver, "type": m_type, "input_size": input_size, "url": m_url} | |
| return m_dict, size_span | |
| def gen_dataframe(url="https://pytorch.org/vision/main/_modules/"): | |
| torch_page = parse_url(url) | |
| article = torch_page.find("article", {"id": "pytorch-article"}) | |
| ul = article.find("ul").find("ul") | |
| in1k_v1, in1k_v2 = [], [] | |
| for li in tqdm(ul.find_all("li"), desc="Crawling cv backbone info..."): | |
| name = str(li.text) | |
| if name.__contains__("torchvision.models.") and len(name.split(".")) == 3: | |
| if ( | |
| name.__contains__("_api") | |
| or name.__contains__("feature_extraction") | |
| or name.__contains__("maxvit") | |
| ): | |
| continue | |
| href = li.find("a").get("href") | |
| model_page = parse_url(url + href) | |
| divs = model_page.select("div.viewcode-block") | |
| for div in divs: | |
| div_id = str(div["id"]) | |
| if div_id.__contains__("_Weights"): | |
| m_ver = div_id.split("_Weight")[0].lower() | |
| if m_ver.__contains__("swin_v2_"): | |
| continue | |
| m_type = special_type(m_ver) | |
| in1k_v1_span = div.find( | |
| name="span", attrs={"class": "n"}, string="IMAGENET1K_V1" | |
| ) | |
| if not in1k_v1_span: | |
| continue | |
| m_dict, size_span = info_on_dataset(m_ver, m_type, in1k_v1_span) | |
| in1k_v1.append(m_dict) | |
| in1k_v2_span = size_span.find_next_sibling( | |
| name="span", attrs={"class": "n"}, string="IMAGENET1K_V2" | |
| ) | |
| if in1k_v2_span: | |
| m_dict, _ = info_on_dataset(m_ver, m_type, in1k_v2_span) | |
| in1k_v2.append(m_dict) | |
| dataset = {"IMAGENET1K_V1": in1k_v1, "IMAGENET1K_V2": in1k_v2} | |
| with open("IMAGENET1K_V1.jsonl", "w", encoding="utf-8") as jsonl_file: | |
| for item in in1k_v1: | |
| jsonl_file.write(json.dumps(item) + "\n") | |
| with open("IMAGENET1K_V2.jsonl", "w", encoding="utf-8") as jsonl_file: | |
| for item in in1k_v2: | |
| jsonl_file.write(json.dumps(item) + "\n") | |
| return dataset | |
| def inference(subset): | |
| cache_json = f"{subset}.jsonl" | |
| if os.path.exists(cache_json): | |
| with open(cache_json, "r", encoding="utf-8") as jsonl_file: | |
| dataset = [json.loads(line) for line in jsonl_file] | |
| else: | |
| dataset = gen_dataframe()[subset] | |
| return pd.DataFrame(dataset), cache_json | |
| def sync(subset): | |
| cache_json = f"{subset}.jsonl" | |
| if os.path.exists(cache_json): | |
| os.remove(cache_json) | |
| return None | |
| with gr.Blocks() as demo: | |
| with gr.Row(): | |
| subset_opt = gr.Dropdown( | |
| choices=["IMAGENET1K_V1", "IMAGENET1K_V2"], value="IMAGENET1K_V1" | |
| ) | |
| sync_btn = gr.Button("Clean cache") | |
| dld_file = gr.components.File(label="Download JSON lines") | |
| with gr.Row(): | |
| data_frame = gr.Dataframe(headers=["ver", "type", "input_size", "url"]) | |
| subset_opt.change(inference, inputs=subset_opt, outputs=[data_frame, dld_file]) | |
| sync_btn.click(sync, inputs=subset_opt, outputs=dld_file) | |
| demo.launch() | |