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