|
from typing import List |
|
|
|
import gradio as gr |
|
import supervision as sv |
|
import torch |
|
from PIL import Image |
|
from ultralytics import YOLO |
|
|
|
MARKDOWN = """ |
|
# Orang Outan Detection |
|
""" |
|
EXAMPLES = [] |
|
|
|
DEVICE = "cuda" if torch.cuda.is_available() else "cpu" |
|
|
|
YOLO_MODEL = YOLO("train_7best.pt") |
|
|
|
BOX_ANNOTATOR = sv.BoxAnnotator() |
|
|
|
|
|
def annotate( |
|
image_bgr_numpy: Image.Image, |
|
detections: sv.Detections, |
|
annotator: sv.BoxAnnotator, |
|
labels: str, |
|
) -> Image.Image: |
|
annotated_bgr_image = annotator.annotate( |
|
scene=image_bgr_numpy, detections=detections, labels=labels |
|
) |
|
return Image.fromarray(annotated_bgr_image[:, :, ::-1]) |
|
|
|
|
|
def inference(image_rgb_pil: Image.Image, confidence: float) -> List[Image.Image]: |
|
output = YOLO_MODEL(image_rgb_pil, verbose=False)[0] |
|
detections = sv.Detections.from_ultralytics(output) |
|
|
|
detections = detections[detections.confidence >= confidence] |
|
|
|
labels = [ |
|
f"{output.names[class_id]} {confidence:0.2f}" |
|
for _, _, confidence, class_id, _ in detections |
|
] |
|
|
|
return annotate( |
|
image_bgr_numpy=output.orig_img.copy(), |
|
detections=detections, |
|
annotator=BOX_ANNOTATOR, |
|
labels=labels, |
|
) |
|
|
|
|
|
def run_demo(): |
|
custom_theme = gr.themes.Soft(primary_hue="blue").set( |
|
button_secondary_background_fill="*neutral_100", |
|
button_secondary_background_fill_hover="*neutral_200", |
|
) |
|
|
|
with gr.Blocks(theme=custom_theme, css="style.css") as demo: |
|
gr.Markdown(MARKDOWN) |
|
|
|
with gr.Row(): |
|
with gr.Column(): |
|
input_image = gr.Image(image_mode="RGB", type="pil", height=500) |
|
confidence_slider = gr.Slider( |
|
label="Confidence", minimum=0.1, maximum=1.0, step=0.05, value=0.6 |
|
) |
|
submit_button = gr.Button("Submit") |
|
output_image = gr.Image(label="Results", type="pil") |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
submit_button.click( |
|
inference, |
|
inputs=[input_image, confidence_slider], |
|
outputs=output_image, |
|
queue=True, |
|
) |
|
demo.queue(max_size=20, api_open=False).launch() |
|
|
|
|
|
if __name__ == "__main__": |
|
run_demo() |
|
|
