Spaces:
Runtime error
Runtime error
| import torch | |
| import gradio as gr | |
| import cv2 | |
| import numpy as np | |
| import random | |
| import numpy as np | |
| from models.experimental import attempt_load | |
| from utils.general import check_img_size, non_max_suppression, \ | |
| scale_coords | |
| from utils.plots import plot_one_box | |
| from utils.torch_utils import time_synchronized | |
| import time | |
| from ultralytics import YOLO | |
| from track import MOT | |
| 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 = ["animal", | |
| "autorickshaw", | |
| "bicycle", | |
| "bus", | |
| "car", | |
| "motorcycle", | |
| "person", | |
| "rider", | |
| "traffic light", | |
| "traffic sign", | |
| "truck" | |
| ] | |
| #colors = [[random.randint(0, 255) for _ in range(3)] for _ in names] | |
| colors = { | |
| "animal": [246,198, 145], | |
| "autorickshaw": [255,204, 54], | |
| "bicycle": [119,11, 32], | |
| "bus": [ 0,60,100], | |
| "car": [ 0,0,142], | |
| "motorcycle": [ 0,0,230], | |
| "person": [220,20, 60], | |
| "rider": [255,0, 0], | |
| "traffic light": [250,170, 30], | |
| "traffic sign": [220,220, 0], | |
| "truck": [ 0,0, 70] | |
| } | |
| def detectv7(img,model,device,iou_threshold=0.45,confidence_threshold=0.25): | |
| imgsz = 640 | |
| img = np.array(img) | |
| stride = int(model.stride.max()) # model stride | |
| imgsz = check_img_size(imgsz, s=stride) # check img_size | |
| # Get names and colors | |
| names = model.module.names if hasattr(model, 'module') else model.names | |
| # Run inference | |
| imgs = img.copy() # for NMS | |
| image, ratio, dwdh = letterbox(img, auto=False) | |
| 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 | |
| t1 = time_synchronized() | |
| start = time.time() | |
| with torch.no_grad(): # Calculating gradients would cause a GPU memory leak | |
| pred = model(img,augment=True)[0] | |
| fps_inference = 1/(time.time()-start) | |
| t2 = time_synchronized() | |
| # Apply NMS | |
| pred = non_max_suppression(pred, confidence_threshold, iou_threshold, classes=None, agnostic=True) | |
| t3 = time_synchronized() | |
| for i, det in enumerate(pred): # detections per image | |
| if len(det): | |
| # Rescale boxes from img_size to im0 size | |
| det[:, :4] = scale_coords(img.shape[2:], det[:, :4], imgs.shape).round() | |
| # Write results | |
| for *xyxy, conf, cls in reversed(det): | |
| label = f'{names[int(cls)]} {conf:.2f}' | |
| plot_one_box(xyxy, imgs, label=label, color=colors[names[int(cls)]], line_thickness=1) | |
| return imgs,fps_inference | |
| def detectv8(img,model,device,iou_threshold=0.45,confidence_threshold=0.25): | |
| img = np.array(img) | |
| # Inference | |
| t1 = time_synchronized() | |
| start = time.time() | |
| results= model.predict(img,conf=confidence_threshold, iou=iou_threshold) | |
| fps_inference = 1/(time.time()-start) | |
| if torch.cuda.is_available(): | |
| boxes= results[0].boxes.cpu().numpy() | |
| else: | |
| boxes=results[0].boxes.numpy() | |
| for bbox in boxes: | |
| #print(f'{colors[names[int(bbox.cls[0])]]}') | |
| label = f'{names[int(bbox.cls[0])]} {bbox.conf[0]:.2f}' | |
| plot_one_box(bbox.xyxy[0],img,colors[names[int(bbox.cls[0])]],label, line_thickness=1) | |
| return img,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 | |
| model_path = 'weights/'+str(model_link)+'.pt' | |
| if model_link== 'yolov8m': | |
| model = YOLO(model_path) | |
| return detectv8(img,model,device,iou_threshold,confidence_threshold) | |
| else: | |
| model = attempt_load(model_path, map_location=device) | |
| return detectv7(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_path = 'weights/'+str(model_link)+'.pt' | |
| if model_link== 'yolov8m': | |
| model = YOLO(model_path) | |
| else: | |
| model = attempt_load(model_path, map_location=device) | |
| 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 | |
| if model_link== 'yolov8m': | |
| frame,fps = detectv8(frame,model,device,iou_threshold,confidence_threshold) | |
| else: | |
| frame,fps = detectv7(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) | |
| def inference_comp(image,iou_threshold,confidence_threshold): | |
| v8_out, v8_fps = inference(image, "yolov8m",iou_threshold,confidence_threshold) | |
| v7_out, v7_fps = inference(image, "yolov7",iou_threshold,confidence_threshold) | |
| return v7_out,v8_out,v7_fps,v8_fps | |
| def MODT(sourceVideo, trackingmethod): | |
| #model_path = 'weights/'+str(model_link)+'.pt' | |
| model_path = 'weights/yolov8m.pt' | |
| return MOT(model_path, trackingmethod, sourceVideo), 30 | |
| examples_images = ['data/images/1.jpg', | |
| 'data/images/2.jpg', | |
| 'data/images/bus.jpg', | |
| 'data/images/3.jpg'] | |
| examples_videos = ['data/video/1.mp4','data/video/2.mp4'] | |
| models = ['yolov8m','yolov7','yolov7t'] | |
| trackers = ['strongsort', 'bytetrack', 'ocsort'] | |
| with gr.Blocks() as demo: | |
| gr.Markdown("## IDD Inference on Yolo V7 and V8 ") | |
| with gr.Tab("Image"): | |
| gr.Markdown("## Yolo V7 and V8 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(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.5) | |
| image_conf_threshold = gr.Slider(label="Confidence Threshold",interactive=True, minimum=0.0, maximum=1.0, value=0.6) | |
| gr.Examples(examples=examples_images,inputs=image_input,outputs=image_output) | |
| text_button = gr.Button("Detect") | |
| with gr.Tab("Video"): | |
| gr.Markdown("## Yolo V7 and V8 Inference on Video") | |
| with gr.Row(): | |
| video_input = gr.Video(type='pil', label="Input Video", source="upload") | |
| video_output = gr.Video(type="pil", label="Output Video",format="mp4") | |
| fps_video = gr.Number(0,label='FPS') | |
| video_drop = gr.Dropdown(label="Model", choices=models,value=models[0]) | |
| video_iou_threshold = gr.Slider(label="IOU Threshold",interactive=True, minimum=0.0, maximum=1.0, value=0.5) | |
| video_conf_threshold = gr.Slider(label="Confidence Threshold",interactive=True, minimum=0.0, maximum=1.0, value=0.6) | |
| gr.Examples(examples=examples_videos,inputs=video_input,outputs=video_output) | |
| with gr.Row(): | |
| video_button_detect = gr.Button("Detect") | |
| with gr.Tab("Compare Models"): | |
| gr.Markdown("## YOLOv7 vs YOLOv8 Object detection comparision") | |
| with gr.Row(): | |
| image_comp_input = gr.Image(type='pil', label="Input Image", source="upload") | |
| with gr.Row(): | |
| image_comp_iou_threshold = gr.Slider(label="IOU Threshold",interactive=True, minimum=0.0, maximum=1.0, value=0.5) | |
| image_comp_conf_threshold = gr.Slider(label="Confidence Threshold",interactive=True, minimum=0.0, maximum=1.0, value=0.6) | |
| text_comp_button = gr.Button("Detect") | |
| with gr.Row(): | |
| image_comp_output_v7 = gr.Image(type='pil', label="YOLOv7 Output Image", source="upload") | |
| image_comp_output_v8 = gr.Image(type='pil', label="YOLOv8 Output Image", source="upload") | |
| with gr.Row(): | |
| v7_fps_image = gr.Number(0,label='v7 FPS') | |
| v8_fps_image = gr.Number(0,label='v8 FPS') | |
| gr.Examples(examples=examples_images,inputs=image_comp_input,outputs=[image_comp_output_v7,image_comp_output_v8]) | |
| with gr.Tab("Video Tacking"): | |
| gr.Markdown("## MOT using YoloV8 detection with tracking") | |
| with gr.Row(): | |
| videotr_input = gr.Video(type='pil', label="Input Video", source="upload") | |
| videotr_output = gr.Video(type="pil", label="Output Video",format="mp4") | |
| fpstr_video = gr.Number(0,label='FPS') | |
| tracking_drop = gr.Dropdown(choices=trackers,value=trackers[0], label="Select the tracking method") | |
| videotr_iou_threshold = gr.Slider(label="IOU Threshold",interactive=True, minimum=0.0, maximum=1.0, value=0.5) | |
| videotr_conf_threshold = gr.Slider(label="Confidence Threshold",interactive=True, minimum=0.0, maximum=1.0, value=0.6) | |
| gr.Examples(examples=examples_videos,inputs=videotr_input,outputs=videotr_output) | |
| video_button_track = gr.Button("Track") | |
| text_button.click(inference, inputs=[image_input,image_drop, | |
| image_iou_threshold,image_conf_threshold], | |
| outputs=[image_output,fps_image]) | |
| video_button_detect.click(inference2, inputs=[video_input,video_drop, | |
| video_iou_threshold,video_conf_threshold], | |
| outputs=[video_output,fps_video]) | |
| text_comp_button.click(inference_comp,inputs=[image_comp_input, | |
| image_comp_iou_threshold, | |
| image_comp_conf_threshold], | |
| outputs=[image_comp_output_v7,image_comp_output_v8,v7_fps_image,v8_fps_image]) | |
| video_button_track.click(MODT,inputs=[videotr_input, tracking_drop], | |
| outputs=[videotr_output, fpstr_video]) | |
| demo.launch(debug=True,enable_queue=True) |