Upload 2 files

app.py

@@ -4,7 +4,7 @@ import gradio as gr
 from datasets import load_dataset
 from ultralyticsplus import YOLO, render_result, postprocess_classify_output
 
-from utils import load_models_from_txt_files
+from utils import load_models_from_txt_files, get_dataset_id_from_model_id, get_task_from_readme
 
 EXAMPLE_IMAGE_DIR = 'example_images'
 
@@ -17,6 +17,7 @@ DEFAULT_CLS_DATASET_ID = 'keremberke/chest-xray-classification'
 
 # load model ids and default models
 det_model_ids, seg_model_ids, cls_model_ids = load_models_from_txt_files()
+task_to_model_ids = {'detect': det_model_ids, 'segment': seg_model_ids, 'classify': cls_model_ids}
 det_model = YOLO(DEFAULT_DET_MODEL_ID)
 det_model_id = DEFAULT_DET_MODEL_ID
 seg_model = YOLO(DEFAULT_SEG_MODEL_ID)
@@ -25,22 +26,25 @@ cls_model = YOLO(DEFAULT_CLS_MODEL_ID)
 cls_model_id = DEFAULT_CLS_MODEL_ID
 
 
-def get_examples(model_id, dataset_id, task):
+def get_examples(task):
     examples = []
-    ds = load_dataset(dataset_id, name="mini")["validation"]
     Path(EXAMPLE_IMAGE_DIR).mkdir(parents=True, exist_ok=True)
-    for ind in range(min(5, len(ds))):
-        jpeg_image_file = ds[ind]["image"]
-        image_file_path = str(Path(EXAMPLE_IMAGE_DIR) / f"{task}_example_{ind}.jpg")
-        jpeg_image_file.save(image_file_path, format='JPEG', quality=100)
-        image_path = os.path.abspath(image_file_path)
-        examples.append([image_path, model_id, 0.25])
+    for model_id in task_to_model_ids[task]:
+        dataset_id = get_dataset_id_from_model_id(model_id)
+        ds = load_dataset(dataset_id, name="mini")["validation"]
+        for ind in range(min(2, len(ds))):
+            jpeg_image_file = ds[ind]["image"]
+            image_file_path = str(Path(EXAMPLE_IMAGE_DIR) / f"{task}_example_{ind}.jpg")
+            jpeg_image_file.save(image_file_path, format='JPEG', quality=100)
+            image_path = os.path.abspath(image_file_path)
+            examples.append([image_path, model_id, 0.25])
     return examples
 
+
 # load default examples using default datasets
-det_examples = get_examples(DEFAULT_DET_MODEL_ID, DEFAULT_DET_DATASET_ID, 'detect')
-seg_examples = get_examples(DEFAULT_SEG_MODEL_ID, DEFAULT_SEG_DATASET_ID, 'segment')
-cls_examples = get_examples(DEFAULT_CLS_MODEL_ID, DEFAULT_CLS_DATASET_ID, 'classification')
+det_examples = get_examples('detect')
+seg_examples = get_examples('segment')
+cls_examples = get_examples('classify')
 
 
 def predict(image, model_id, threshold):
@@ -120,12 +124,12 @@ with gr.Blocks() as demo:
             with gr.Column():
                 detect_output = gr.Image(label="Predictions:", interactive=False)
         with gr.Row():
-            gr.Examples(
+            detect_examples = gr.Examples(
                     det_examples,
                     inputs=[detect_input, detect_model_id, detect_threshold],
                     outputs=detect_output,
                     fn=predict,
-                    cache_examples=True,
+                    cache_examples=False,
                 )
     with gr.Tab("Segmentation"):
         with gr.Row():
@@ -137,7 +141,7 @@ with gr.Blocks() as demo:
             with gr.Column():
                 segment_output = gr.Image(label="Predictions:", interactive=False)
         with gr.Row():
-            gr.Examples(
+            segment_examples = gr.Examples(
                     seg_examples,
                     inputs=[segment_input, segment_model_id, segment_threshold],
                     outputs=segment_output,
@@ -156,7 +160,7 @@ with gr.Blocks() as demo:
                     label="Predictions:", show_label=True, num_top_classes=5
                 )
         with gr.Row():
-            gr.Examples(
+            classify_examples = gr.Examples(
                     cls_examples,
                     inputs=[classify_input, classify_model_id, classify_threshold],
                     outputs=classify_output,

utils.py

@@ -1,3 +1,7 @@
+import requests
+import re
+
+
 DET_MODELS_FILENAME = 'det_models.txt'
 SEG_MODELS_FILENAME = 'seg_models.txt'
 CLS_MODELS_FILENAME = 'cls_models.txt'
@@ -11,4 +15,62 @@ def load_models_from_txt_files():
         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
+    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