File size: 2,820 Bytes
1a2f039 81dce6e 1a2f039 81dce6e 1a2f039 81dce6e 1a2f039 db5fd2f 1a2f039 25038c9 1a2f039 8196cbc 1a2f039 81dce6e aa09faf 1a2f039 2f10a95 |
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 |
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"), # Accepts image input
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()
|