Spaces:
Runtime error
Runtime error
File size: 8,907 Bytes
307beef 10f3130 307beef 10f3130 307beef 10f3130 307beef 651391d 10f3130 651391d 10f3130 651391d 307beef 10f3130 307beef 10f3130 307beef 651391d 307beef 651391d 307beef 651391d 307beef |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 |
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, get_dataset_id_from_model_id, get_task_from_readme
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()
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)
seg_model_id = DEFAULT_SEG_MODEL_ID
cls_model = YOLO(DEFAULT_CLS_MODEL_ID)
cls_model_id = DEFAULT_CLS_MODEL_ID
def get_examples(task):
examples = []
Path(EXAMPLE_IMAGE_DIR).mkdir(parents=True, exist_ok=True)
image_ind = 0
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_{image_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])
image_ind += 1
return examples
# load default examples using default datasets
det_examples = get_examples('detect')
seg_examples = get_examples('segment')
cls_examples = get_examples('classify')
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():
half_ind = int(len(det_examples) / 2)
with gr.Column():
detect_examples = gr.Examples(
det_examples[:half_ind],
inputs=[detect_input, detect_model_id, detect_threshold],
outputs=detect_output,
fn=predict,
cache_examples=False,
)
with gr.Column():
detect_examples = gr.Examples(
det_examples[:half_ind],
inputs=[detect_input, detect_model_id, detect_threshold],
outputs=detect_output,
fn=predict,
cache_examples=False,
)
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():
half_ind = int(len(det_examples) / 2)
with gr.Column():
segment_examples = gr.Examples(
seg_examples[:half_ind],
inputs=[segment_input, segment_model_id, segment_threshold],
outputs=segment_output,
fn=predict,
cache_examples=False,
)
with gr.Column():
segment_examples = gr.Examples(
seg_examples[:half_ind],
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():
half_ind = int(len(det_examples) / 2)
with gr.Column():
classify_examples = gr.Examples(
cls_examples[half_ind:],
inputs=[classify_input, classify_model_id, classify_threshold],
outputs=classify_output,
fn=predict,
cache_examples=False,
)
with gr.Column():
classify_examples = gr.Examples(
cls_examples[:half_ind],
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) |