|
|
from yolox_onnx import YOLOX_ONNX |
|
|
import albumentations as A |
|
|
import gradio as gr |
|
|
import cv2 |
|
|
import numpy as np |
|
|
|
|
|
def show_example(path): |
|
|
return cv2.cvtColor(cv2.imread(path), cv2.COLOR_BGR2RGB) |
|
|
|
|
|
|
|
|
def get_response(input_img,add_snow,add_rain,add_fog,confidence_threshold,iou_threshold): |
|
|
|
|
|
''' |
|
|
detects all possible pedestrians in the image and recognizes it |
|
|
Args: |
|
|
input_img (numpy array): one image of type numpy array |
|
|
add_snow (boolean) : apply random snow augmentation |
|
|
add_rain (boolean) : apply random rain augmentation |
|
|
add_fog (boolean) : apply random fog augmentation |
|
|
confidence_threshold (float) : minimum confidence prob required for bounding box candidate |
|
|
iou_threshold (float) : intersection threshold above which bounding box will be neglected |
|
|
|
|
|
Returns: |
|
|
return img(numpy array): image with bounding boxes of pedestrians |
|
|
''' |
|
|
|
|
|
|
|
|
if not hasattr(input_img,'shape'): |
|
|
return "invalid input",input_img |
|
|
|
|
|
input_img=cv2.cvtColor(input_img,cv2.COLOR_RGB2BGR).astype(np.float32) |
|
|
|
|
|
pedestrian_detector.predict(input_img,confidence_threshold,iou_threshold) |
|
|
out_img=pedestrian_detector.output_img.astype(np.uint8) |
|
|
|
|
|
if add_snow: |
|
|
out_img = weather_transform["add_snow"] (image=out_img)['image'] |
|
|
elif add_rain: |
|
|
out_img = weather_transform["add_rain"](image=out_img)['image'] |
|
|
elif add_fog: |
|
|
out_img = weather_transform["add_fog"](image=out_img)['image'] |
|
|
else: |
|
|
pass |
|
|
|
|
|
|
|
|
return cv2.cvtColor(out_img,cv2.COLOR_BGR2RGB).astype(np.uint8) |
|
|
|
|
|
|
|
|
if __name__ == "__main__": |
|
|
weather_transform = {} |
|
|
weather_transform["add_rain"] = A.RandomRain(brightness_coefficient=0.9, drop_width=1, blur_value=5, p=1) |
|
|
weather_transform["add_snow"] = A.RandomSnow(brightness_coeff=2.5, snow_point_lower=0.3, snow_point_upper=0.5, p=1) |
|
|
weather_transform["add_fog"] = A.RandomFog(fog_coef_lower=0.7, fog_coef_upper=0.8, alpha_coef=0.1, p=1) |
|
|
|
|
|
|
|
|
pedestrian_detector=YOLOX_ONNX('models/pedestrian-detection-best95.onnx') |
|
|
iface = gr.Interface( |
|
|
cache_examples=False, |
|
|
fn=get_response, |
|
|
inputs=[gr.Image(type="numpy"), |
|
|
gr.Checkbox(label="add_snow"), |
|
|
gr.Checkbox(label="add_rain"), |
|
|
gr.Checkbox(label="add_fog"), |
|
|
gr.Slider(0, 1,value=0.5, step=0.01, label="confidence_threshold"), |
|
|
gr.Slider(0, 1,value=0.4, step=0.01, label="iou_threshold")], |
|
|
examples=[[show_example('test-images/test1.jpg')],[show_example('test-images/test2.jpg')],[show_example('test-images/test3.jpg')]], |
|
|
outputs=[gr.Image(type="numpy")], |
|
|
title="Pedestrian Detection with All weather augmentation", |
|
|
description="Upload images for pedestrian detection") |
|
|
|
|
|
iface.launch() |
|
|
|
