import torch from models.common import DetectMultiBackend from utils.general import (check_img_size, cv2, non_max_suppression, scale_boxes) from utils.plots import Annotator, colors import numpy as np import gradio as gr import time data = 'data/coco128.yaml' def letterbox(im, new_shape=(640, 640), color=(114, 114, 114), auto=True, scaleup=True, stride=32): # Resize and pad image while meeting stride-multiple constraints shape = im.shape[:2] # current shape [height, width] if isinstance(new_shape, int): new_shape = (new_shape, new_shape) # Scale ratio (new / old) r = min(new_shape[0] / shape[0], new_shape[1] / shape[1]) if not scaleup: # only scale down, do not scale up (for better val mAP) r = min(r, 1.0) # Compute padding new_unpad = int(round(shape[1] * r)), int(round(shape[0] * r)) dw, dh = new_shape[1] - new_unpad[0], new_shape[0] - new_unpad[1] # wh padding if auto: # minimum rectangle dw, dh = np.mod(dw, stride), np.mod(dh, stride) # wh padding dw /= 2 # divide padding into 2 sides dh /= 2 if shape[::-1] != new_unpad: # resize im = cv2.resize(im, new_unpad, interpolation=cv2.INTER_LINEAR) top, bottom = int(round(dh - 0.1)), int(round(dh + 0.1)) left, right = int(round(dw - 0.1)), int(round(dw + 0.1)) im = cv2.copyMakeBorder(im, top, bottom, left, right, cv2.BORDER_CONSTANT, value=color) # add border return im, r, (dw, dh) names = ['person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus', 'train', 'truck', 'boat', 'traffic light', 'fire hydrant', 'stop sign', 'parking meter', 'bench', 'bird', 'cat', 'dog', 'horse', 'sheep', 'cow', 'elephant', 'bear', 'zebra', 'giraffe', 'backpack', 'umbrella', 'handbag', 'tie', 'suitcase', 'frisbee', 'skis', 'snowboard', 'sports ball', 'kite', 'baseball bat', 'baseball glove', 'skateboard', 'surfboard', 'tennis racket', 'bottle', 'wine glass', 'cup', 'fork', 'knife', 'spoon', 'bowl', 'banana', 'apple', 'sandwich', 'orange', 'broccoli', 'carrot', 'hot dog', 'pizza', 'donut', 'cake', 'chair', 'couch', 'potted plant', 'bed', 'dining table', 'toilet', 'tv', 'laptop', 'mouse', 'remote', 'keyboard', 'cell phone', 'microwave', 'oven', 'toaster', 'sink', 'refrigerator', 'book', 'clock', 'vase', 'scissors', 'teddy bear', 'hair drier', 'toothbrush'] def detect(im,model,device,iou_threshold=0.45,confidence_threshold=0.25): im = np.array(im) imgsz=(640, 640) # inference size (pixels) data = 'data/coco128.yaml' # data.yaml path # Load model stride, names, pt = model.stride, model.names, model.pt imgsz = check_img_size(imgsz, s=stride) # check image size # Run inference # model.warmup(imgsz=(1)) # warmup imgs = im.copy() # for NMS image, ratio, dwdh = letterbox(im, auto=False) print(image.shape) image = image.transpose((2, 0, 1)) img = torch.from_numpy(image).to(device) img = img.float() # uint8 to fp16/32 img /= 255.0 # 0 - 255 to 0.0 - 1.0 if img.ndimension() == 3: img = img.unsqueeze(0) # Inference start = time.time() pred = model(img, augment=False) fps_inference = 1/(time.time()-start) # NMS pred = non_max_suppression(pred, confidence_threshold, iou_threshold, None, False, max_det=10) for i, det in enumerate(pred): # detections per image if len(det): # Rescale boxes from img_size to im0 size det[:, :4] = scale_boxes(img.shape[2:], det[:, :4], imgs.shape).round() annotator = Annotator(imgs, line_width=3, example=str(names)) hide_labels = False hide_conf = False # Write results for *xyxy, conf, cls in reversed(det): c = int(cls) # integer class label = None if hide_labels else (names[c] if hide_conf else f'{names[c]} {conf:.2f}') print(xyxy,label) annotator.box_label(xyxy, label, color=colors(c, True)) return imgs,fps_inference def inference(img,model_link,iou_threshold,confidence_threshold): print(model_link) device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu') # Load model device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu') model = DetectMultiBackend('weights/'+str(model_link)+'.pt', device=device, dnn=False, data=data, fp16=False) return detect(img,model,device,iou_threshold,confidence_threshold) def inference2(video,model_link,iou_threshold,confidence_threshold): print(model_link) device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu') # Load model model = DetectMultiBackend('weights/'+str(model_link)+'.pt', device=device, dnn=False, data=data, fp16=False) frames = cv2.VideoCapture(video) fps = frames.get(cv2.CAP_PROP_FPS) image_size = (int(frames.get(cv2.CAP_PROP_FRAME_WIDTH)),int(frames.get(cv2.CAP_PROP_FRAME_HEIGHT))) finalVideo = cv2.VideoWriter('output.mp4',cv2.VideoWriter_fourcc(*'VP90'), fps, image_size) fps_video = [] while frames.isOpened(): ret,frame = frames.read() if not ret: break frame,fps = detect(frame,model,device,iou_threshold,confidence_threshold) fps_video.append(fps) finalVideo.write(frame) frames.release() finalVideo.release() return 'output.mp4',np.mean(fps_video) examples_images = ['data/images/bus.jpg', 'data/images/zidane.jpg',] examples_videos = ['data/video/input_0.mp4', 'data/video/input_1.mp4'] models = ['yolov5s','yolov5n','yolov5m','yolov5l','yolov5x'] with gr.Blocks() as demo: gr.Markdown("## YOLOv5 Inference") with gr.Tab("Image"): gr.Markdown("## YOLOv5 Inference on Image") with gr.Row(): image_input = gr.Image(type='pil', label="Input Image", source="upload") image_output = gr.Image(type='pil', label="Output Image", source="upload") fps_image = gr.Number(value=0,label='FPS') image_drop = gr.Dropdown(choices=models,value=models[0]) image_iou_threshold = gr.Slider(label="IOU Threshold",interactive=True, minimum=0.0, maximum=1.0, value=0.45) image_conf_threshold = gr.Slider(label="Confidence Threshold",interactive=True, minimum=0.0, maximum=1.0, value=0.25) gr.Examples(examples=examples_images,inputs=image_input,outputs=image_output) text_button = gr.Button("Detect") with gr.Tab("Video"): gr.Markdown("## YOLOv5 Inference on Video") with gr.Row(): video_input = gr.Video(type='pil', label="Input Image", source="upload") video_output = gr.Video(type="pil", label="Output Image",format="mp4") fps_video = gr.Number(value=0,label='FPS') video_drop = gr.Dropdown(choices=models,value=models[0]) video_iou_threshold = gr.Slider(label="IOU Threshold",interactive=True, minimum=0.0, maximum=1.0, value=0.45) video_conf_threshold = gr.Slider(label="Confidence Threshold",interactive=True, minimum=0.0, maximum=1.0, value=0.25) gr.Examples(examples=examples_videos,inputs=video_input,outputs=video_output) video_button = gr.Button("Detect") with gr.Tab("Webcam Video"): gr.Markdown("## YOLOv5 Inference on Webcam Video") gr.Markdown("Coming Soon") text_button.click(inference, inputs=[image_input,image_drop, image_iou_threshold,image_conf_threshold], outputs=[image_output,fps_image]) video_button.click(inference2, inputs=[video_input,video_drop, video_iou_threshold,video_conf_threshold], outputs=[video_output,fps_video]) demo.launch()