app.py

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()