utils.py

import requests
import re


DET_MODELS_FILENAME = 'det_models.txt'
SEG_MODELS_FILENAME = 'seg_models.txt'
CLS_MODELS_FILENAME = 'cls_models.txt'


def load_models_from_txt_files():
    """Load models from txt files."""
    with open(DET_MODELS_FILENAME, 'r') as file:
        det_models = [line.strip() for line in file]
    with open(SEG_MODELS_FILENAME, 'r') as file:
        seg_models = [line.strip() for line in file]
    with open(CLS_MODELS_FILENAME, 'r') as file:
        cls_models = [line.strip() for line in file]
    return det_models, seg_models, cls_models


def get_dataset_id_from_model_id(model_id):
    """
    Gets the dataset ID from the README file for a given Hugging Face model ID.
    
    Args:
        model_id (str): The Hugging Face model ID.
    
    Returns:
        The dataset ID as a string, or None if the dataset ID cannot be found.
    """
    # Define the URL of the README file for the model
    readme_url = f"https://huggingface.co/{model_id}/raw/main/README.md"

    # Make a GET request to the README URL and get the contents
    response = requests.get(readme_url)
    readme_contents = response.text

    # Use regular expressions to search for the dataset ID in the README file
    match = re.search(r"datasets:\s*\n- (\S+)", readme_contents)

    # If a match is found, extract the dataset ID and return it. Otherwise, return None.
    if match is not None:
        dataset_id = match.group(1)
        return dataset_id
    else:
        return None


def get_task_from_readme(model_id):
    """
    Gets the task from the README file for a given Hugging Face model ID.
    
    Args:
        model_id (str): The Hugging Face model ID.
    
    Returns:
        The task as a string ("detect", "segment", or "classify"), or None if the task cannot be found.
    """
    # Define the URL of the README file for the model
    readme_url = f"https://huggingface.co/{model_id}/raw/main/README.md"

    # Make a GET request to the README URL and get the contents
    response = requests.get(readme_url)
    readme_contents = response.text

    # Use regular expressions to search for the task in the tags section of the README file
    if re.search(r"tags:", readme_contents):
        if re.search(r"object-detection", readme_contents):
            return "detect"
        elif re.search(r"image-segmentation", readme_contents):
            return "segment"
        elif re.search(r"image-classification", readme_contents):
            return "classify"

    # If the task cannot be found, return None
    return None