Spaces:
Running
Running
File size: 3,214 Bytes
4ba6fde 62c1b39 4ba6fde 62c1b39 4ba6fde |
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 |
import gradio as gr
import PIL.Image as Image
import spaces
import super_gradients
from tools.tools import py_cpu_nms,get_sub_image,filter_small_fp
import cv2
import numpy as np
import os
from classifiers.MixMatch.mixmatch_classification import mixmatch_classifier_inference
@spaces.GPU
def inference_mega_image_yolonas(img, conf_threshold, iou_threshold,height):
record_list = []
model_dir = './checkpoint/yolonas/height_varient/ckpt_best{}.pth'.format(height)
model = super_gradients.training.models.get('yolo_nas_m',num_classes=1,checkpoint_path=model_dir).cuda()
# mega_image = np.array(img)[:, :, ::-1].copy()
mega_image = img
ratio = 1
bbox_list = []
sub_image_list, coor_list = get_sub_image(mega_image, overlap=0.2, ratio=ratio)
for index, sub_image in enumerate(sub_image_list):
# sub_image = cv2.cvtColor(sub_image, cv2.COLOR_BGR2RGB)
# sub_image = Image.fromarray(sub_image)
images_predictions = model.predict(sub_image)
image_prediction = next(iter(images_predictions))
labels = image_prediction.prediction.labels
confidences = image_prediction.prediction.confidence
bboxes = image_prediction.prediction.bboxes_xyxy
for i in range(len(labels)):
label = labels[i]
confidence = confidences[i]
bbox = bboxes[i]
if confidence > conf_threshold:
bbox_list.append([int(coor_list[index][1]+bbox[0]), int(coor_list[index][0]+bbox[1]),int(coor_list[index][1]+bbox[2]), int(coor_list[index][0]+bbox[3]), confidence])
if (len(bbox_list) != 0):
bbox_list = np.asarray([box for box in bbox_list])
box_idx = py_cpu_nms(bbox_list, iou_threshold)
selected_bbox = bbox_list[box_idx]
selected_bbox = sorted(selected_bbox,key = lambda x: x[4],reverse = True)
mega_image = draw_image(mega_image,selected_bbox)
else:
selected_bbox = []
return mega_image,selected_bbox
def draw_image(img,bboxes):
for box in bboxes:
cv2.rectangle(img, (int(box[0]),int(box[1])), (int(box[2]),int(box[3])), (0,255,0), 3)
cv2.putText(img, 'bird', (int(box[0]), int(box[1])), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 255), 1)
return img
@spaces.GPU
def predict_image(img, conf_threshold, iou_threshold,height):
result_image,bbox_list = inference_mega_image_yolonas(img, conf_threshold, iou_threshold,height)
cla_dict = mixmatch_classifier_inference('./checkpoint/classifier/mixmatch/model_best.pth.tar',result_image,bbox_list)
return result_image,cla_dict
iface = gr.Interface(
fn=predict_image,
inputs=[
gr.Image(type="numpy", label="Upload Image"),
gr.Slider(minimum=0, maximum=1, value=0.7, label="Confidence threshold"),
gr.Slider(minimum=0, maximum=1, value=0.3, label="IoU threshold"),
gr.Radio(["15m", "30m", "60m", "90m"], value="15m", label="Height", info="The image taken height"),
],
outputs=[
gr.Image(type="numpy", label="Result"),
gr.Image(type="numpy", label="BarChart")
],
title="Waterfowl detection with YOLONAS",
description="Upload images for Waterfowl object detection.",
)
iface.launch(share=True) |