|
import gradio as gr |
|
import cv2 |
|
import requests |
|
import os |
|
|
|
from ultralytics import YOLO |
|
|
|
model = YOLO('best.pt') |
|
|
|
def normalize_preds(label): |
|
|
|
if label =="Blank": |
|
label = "Deer" |
|
|
|
return label |
|
|
|
def show_preds_image(image): |
|
|
|
image_path = "uploaded_image.jpg" |
|
cv2.imwrite(image_path, image[:, :, ::-1]) |
|
|
|
image = cv2.imread(image_path) |
|
outputs = model.predict(source=image_path) |
|
|
|
names = outputs[0].names |
|
results = outputs[0].boxes.xyxy.cpu().numpy() |
|
probs = outputs[0].boxes.conf.cpu().numpy() |
|
species = outputs[0].boxes.cls.cpu().numpy() |
|
|
|
for idx in range(len(results)): |
|
x1, y1, x2, y2 = results[idx][:4] |
|
|
|
label = normalize_preds(names[int(species[idx])]) |
|
prob_cls = probs[idx] |
|
cv2.rectangle( |
|
image, |
|
(int(x1), int(y1)), |
|
(int(x2), int(y2)), |
|
color=(0, 0, 255), |
|
thickness=4, |
|
lineType=cv2.LINE_AA, |
|
) |
|
cv2.putText( |
|
image, |
|
label, |
|
(int(x1), int(y1) + 20), |
|
cv2.FONT_HERSHEY_SIMPLEX, |
|
0.9, |
|
(0, 0, 255), |
|
2, |
|
cv2.LINE_AA, |
|
) |
|
os.remove(image_path) |
|
return cv2.cvtColor(image, cv2.COLOR_BGR2RGB) |
|
|
|
inputs_image = [ |
|
gr.inputs.Image(label="Upload Image"), |
|
] |
|
outputs_image = [ |
|
gr.outputs.Image(type="numpy"), |
|
] |
|
interface_image = gr.Interface( |
|
fn=show_preds_image, |
|
inputs=inputs_image, |
|
outputs=outputs_image, |
|
title="ReWilding Europe Project - Demo[Yolov8]", |
|
examples=[], |
|
cache_examples=False, |
|
) |
|
|
|
interface_image.launch() |
|
|
