Added yolo v8 model

app.py

@@ -10,8 +10,7 @@ from utils.general import check_img_size, non_max_suppression, \
 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
@@ -71,7 +70,7 @@ colors  = {
 }
 
 
-def detect(img,model,device,iou_threshold=0.45,confidence_threshold=0.25):   
+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
@@ -117,13 +116,33 @@ def detect(img,model,device,iou_threshold=0.45,confidence_threshold=0.25):
 
     return imgs,fps_inference
 
+def detectv8(img,model,device,iou_threshold=0.45,confidence_threshold=0.25):   
+
+    # 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'
-    model = attempt_load(model_path, map_location=device) 
-    return detect(img,model,device,iou_threshold,confidence_threshold)
+    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):
@@ -131,7 +150,10 @@ def inference2(video,model_link,iou_threshold,confidence_threshold):
     device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
     # Load model
     model_path = 'weights/'+str(model_link)+'.pt'
-    model = attempt_load(model_path, map_location=device) 
+    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)))
@@ -141,7 +163,10 @@ def inference2(video,model_link,iou_threshold,confidence_threshold):
         ret,frame = frames.read()
         if not ret:
             break
-        frame,fps = detect(frame,model,device,iou_threshold,confidence_threshold)
+        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()
@@ -156,12 +181,12 @@ examples_images = ['data/images/1.jpg',
             'data/images/3.jpg']
 examples_videos = ['data/video/1.mp4','data/video/2.mp4'] 
 
-models = ['yolov7','yolov7t']
+models = ['yolov8m','yolov7','yolov7t']
 
 with gr.Blocks() as demo:
-    gr.Markdown("## YOLOv7 Inference")
+    gr.Markdown("## IDD Inference on Yolo V7 and V8 ")
     with gr.Tab("Image"):
-        gr.Markdown("## YOLOv7 Inference on 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")
@@ -172,7 +197,7 @@ with gr.Blocks() as demo:
         gr.Examples(examples=examples_images,inputs=image_input,outputs=image_output)
         text_button = gr.Button("Detect")
     with gr.Tab("Video"):
-        gr.Markdown("## YOLOv7 Inference on 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")

requirements.txt

@@ -7,4 +7,5 @@ tqdm>=4.64.0
 seaborn>=0.11.0
 scipy>=1.4.1
 Pillow>=7.1.2
-huggingface-hub >= 0.11.0
+huggingface-hub >= 0.11.0
+ultralytics >=8.0.34

weights/yolov8m.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:54583fba52c0a7551abcd102f8f9652d196ca75e951b7b8a648f31eef7c8a7a9
+size 52018656