ultralytics / app.py
Jason Adrian
initial commit
b7303c6
raw
history blame
2.81 kB
import gradio as gr
import torch
from torchvision.transforms import transforms
import numpy as np
from typing import Optional
import torch.nn as nn
import os
from utils import page_utils
from ultralytics import YOLO
# Load a model
model = YOLO('model_- 14 december 2023 12_01.pt') # pretrained YOLOv8n model
class_names = ['abdominal', 'adult', 'others', 'pediatric', 'spine']
class_names.sort()
examples_dir = "samples"
def image_classifier(inp):
"""Image Classifier Function.
Parameters
----------
inp: Optional[np.ndarray] = None
Input image from callback
Returns
-------
Dict
A dictionary class names and its probability
"""
# If input not valid, return dummy data or raise error
if inp is None:
return {'cat': 0.3, 'dog': 0.7}
result = model(inp)
# postprocess
labeled_result = {class_names[label]: confidence for label, confidence in zip(result.top5, result.top5conf)}
return labeled_result
# gradio code block for input and output
with gr.Blocks() as app:
gr.Markdown("# Lung Cancer Classification")
with open('index.html', encoding="utf-8") as f:
description = f.read()
# gradio code block for input and output
with gr.Blocks(theme=gr.themes.Default(primary_hue=page_utils.KALBE_THEME_COLOR, secondary_hue=page_utils.KALBE_THEME_COLOR).set(
button_primary_background_fill="*primary_600",
button_primary_background_fill_hover="*primary_500",
button_primary_text_color="white",
)) as app:
with gr.Column():
gr.HTML(description)
with gr.Row():
with gr.Column():
inp_img = gr.Image()
with gr.Row():
clear_btn = gr.Button(value="Clear")
process_btn = gr.Button(value="Process", variant="primary")
with gr.Column():
out_txt = gr.Label(label="Probabilities", num_top_classes=5)
process_btn.click(image_classifier, inputs=inp_img, outputs=out_txt)
clear_btn.click(lambda:(
gr.update(value=None),
gr.update(value=None)
),
inputs=None,
outputs=[inp_img, out_txt])
gr.Markdown("## Image Examples")
gr.Examples(
examples=[os.path.join(examples_dir, "1.2.840.113564.1921681202.202011100756242032.1203801020003.dcm.jpeg")
],
inputs=inp_img,
outputs=out_txt,
fn=image_classifier,
cache_examples=False,
)
gr.Markdown(line_breaks=True, value='Author: Jason Adrian ([email protected]) <div class="row"><a href="https://github.com/jasonadriann?tab=repositories"><img alt="GitHub" src="https://img.shields.io/badge/Jason%20Adrian-000000?logo=github"> </div>')
# demo = gr.Interface(fn=image_classifier, inputs="image", outputs="label")
app.launch(share=True)