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 | |
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=2) | |
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) | |
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) | |
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'] | |
def update_markdown(value): | |
return f"## {value} Inference on Image" | |
with gr.Blocks() as demo: | |
gr.Markdown("## IDD Inference on Yolo V7 and V8 ") | |
with gr.Tab("Image"): | |
im_md= 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.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") | |
image_drop.change(update_markdown,image_drop,im_md) | |
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(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("## YOLOv7 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(debug=True,enable_queue=True) |