Upload 7 files

README.md

@@ -1,13 +1,11 @@
 ---
-title: Awesome Yolov8 Models
-emoji: ⚡
-colorFrom: blue
-colorTo: red
+title: Awesome YOLOv8 Models
+emoji: 💯
+colorFrom: green
+colorTo: yellow
 sdk: gradio
-sdk_version: 3.19.1
+sdk_version: 3.17.1
 app_file: app.py
-pinned: false
-license: openrail
----
-
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+pinned: true
+license: mit
+---

app.py

@@ -0,0 +1,177 @@
+import os
+from pathlib import Path
+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
+
+EXAMPLE_IMAGE_DIR = 'example_images'
+
+DEFAULT_DET_MODEL_ID = 'keremberke/yolov8m-valorant-detection'
+DEFAULT_DET_DATASET_ID = 'keremberke/valorant-object-detection'
+DEFAULT_SEG_MODEL_ID = 'keremberke/yolov8s-building-segmentation'
+DEFAULT_SEG_DATASET_ID = 'keremberke/satellite-building-segmentation'
+DEFAULT_CLS_MODEL_ID = 'keremberke/yolov8m-chest-xray-classification'
+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()
+det_model = YOLO(DEFAULT_DET_MODEL_ID)
+det_model_id = DEFAULT_DET_MODEL_ID
+seg_model = YOLO(DEFAULT_SEG_MODEL_ID)
+seg_model_id = DEFAULT_SEG_MODEL_ID
+cls_model = YOLO(DEFAULT_CLS_MODEL_ID)
+cls_model_id = DEFAULT_CLS_MODEL_ID
+
+
+def get_examples(model_id, dataset_id, 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])
+    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')
+
+
+def predict(image, model_id, threshold):
+    """Perform inference on image."""
+    # set task
+    if model_id in det_model_ids:
+        task = 'detect'
+    elif model_id in seg_model_ids:
+        task = 'segment'
+    elif model_id in cls_model_ids:
+        task = 'classify'
+    else:
+        raise ValueError(f"Invalid model_id: {model_id}")
+
+    # set model
+    if task == 'detect':
+        global det_model
+        global det_model_id
+        if model_id != det_model_id:
+            det_model = YOLO(model_id)
+            det_model_id = model_id
+        model = det_model
+    elif task == 'segment':
+        global seg_model
+        global seg_model_id
+        if model_id != seg_model_id:
+            seg_model = YOLO(model_id)
+            seg_model_id = model_id
+        model = seg_model
+    elif task == 'classify':
+        global cls_model
+        global cls_model_id
+        if model_id != cls_model_id:
+            cls_model = YOLO(model_id)
+            cls_model_id = model_id
+        model = cls_model
+    else:
+        raise ValueError(f"Invalid task: {task}")
+
+    # set model parameters
+    model.overrides['conf'] = threshold
+
+    # perform inference
+    results = model.predict(image)
+    print(model_id)
+    print(task)
+
+    if task in ['detect', 'segment']:
+        # draw predictions
+        output = render_result(model=model, image=image, result=results[0])
+    elif task == 'classify':
+        # postprocess classification output
+        output = postprocess_classify_output(model, result=results[0])
+    else:
+        raise ValueError(f"Invalid task: {task}")
+
+    return output
+
+
+with gr.Blocks() as demo:
+    gr.Markdown("""# <p align='center'><img width='500px' src='https://user-images.githubusercontent.com/34196005/215836968-fb54e066-a524-4caf-b469-92bbaa96f921.gif' /></p>
+    <p style='text-align: center'>
+        <br> <a href='https://yolov8.xyz' target='_blank'>project website</a> | <a href='https://github.com/keremberke/awesome-yolov8-models' target='_blank'>project github</a> 
+    </p>
+    <p style='text-align: center'>
+        Follow me for more! 
+        <br> <a href='https://twitter.com/_keremberke' target='_blank'>twitter</a> | <a href='https://github.com/keremberke' target='_blank'>github</a> | <a href='https://www.linkedin.com/in/kerem-berke-bba6a5204/' target='_blank'>linkedin</a> 
+    </p>
+    """)
+    with gr.Tab("Detection"):
+        with gr.Row():
+            with gr.Column():
+                detect_input = gr.Image()
+                detect_model_id = gr.Dropdown(choices=det_model_ids, label="Model:", value=DEFAULT_DET_MODEL_ID, interactive=True)
+                detect_threshold = gr.Slider(maximum=1, step=0.01, value=0.25, label="Threshold:", interactive=True)
+                detect_button = gr.Button("Detect!")
+            with gr.Column():
+                detect_output = gr.Image(label="Predictions:", interactive=False)
+        with gr.Row():
+            gr.Examples(
+                    det_examples,
+                    inputs=[detect_input, detect_model_id, detect_threshold],
+                    outputs=detect_output,
+                    fn=predict,
+                    cache_examples=True,
+                )
+    with gr.Tab("Segmentation"):
+        with gr.Row():
+            with gr.Column():
+                segment_input = gr.Image()
+                segment_model_id = gr.Dropdown(choices=seg_model_ids, label="Model:", value=DEFAULT_SEG_MODEL_ID, interactive=True)
+                segment_threshold = gr.Slider(maximum=1, step=0.01, value=0.25, label="Threshold:", interactive=True)
+                segment_button = gr.Button("Segment!")
+            with gr.Column():
+                segment_output = gr.Image(label="Predictions:", interactive=False)
+        with gr.Row():
+            gr.Examples(
+                    seg_examples,
+                    inputs=[segment_input, segment_model_id, segment_threshold],
+                    outputs=segment_output,
+                    fn=predict,
+                    cache_examples=False,
+                )
+    with gr.Tab("Classification"):
+        with gr.Row():
+            with gr.Column():
+                classify_input = gr.Image()
+                classify_model_id = gr.Dropdown(choices=cls_model_ids, label="Model:", value=DEFAULT_CLS_MODEL_ID, interactive=True)
+                classify_threshold = gr.Slider(maximum=1, step=0.01, value=0.25, label="Threshold:", interactive=True)
+                classify_button = gr.Button("Classify!")
+            with gr.Column():
+                classify_output = gr.Label(
+                    label="Predictions:", show_label=True, num_top_classes=5
+                )
+        with gr.Row():
+            gr.Examples(
+                    cls_examples,
+                    inputs=[classify_input, classify_model_id, classify_threshold],
+                    outputs=classify_output,
+                    fn=predict,
+                    cache_examples=False,
+                )
+
+    detect_button.click(
+        predict, inputs=[detect_input, detect_model_id, detect_threshold], outputs=detect_output
+    )
+    segment_button.click(
+        predict, inputs=[segment_input, segment_model_id, segment_threshold], outputs=segment_output
+    )
+    classify_button.click(
+        predict, inputs=[classify_input, classify_model_id, classify_threshold], outputs=classify_output
+    )
+
+demo.launch(server_port=8080)

cls_models.txt

@@ -0,0 +1,2 @@
+keremberke/yolov8m-shoe-classification
+keremberke/yolov8m-chest-xray-classification

det_models.txt

@@ -0,0 +1,8 @@
+keremberke/yolov8m-valorant-detection
+keremberke/yolov8m-csgo-player-detection
+keremberke/yolov8m-forklift-detection
+keremberke/yolov8m-blood-cell-detection
+keremberke/yolov8m-plane-detection
+keremberke/yolov8m-nlf-head-detection
+keremberke/yolov8m-hard-hat-detection
+keremberke/yolov8m-table-extraction

requirements.txt

@@ -0,0 +1,2 @@
+torch
+ultralyticsplus==0.0.25

seg_models.txt

@@ -0,0 +1,3 @@
+keremberke/yolov8m-pcb-defect-segmentation
+keremberke/yolov8s-building-segmentation
+keremberke/yolov8n-pothole-segmentation

utils.py

@@ -0,0 +1,14 @@
+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