└── YOLOv8InferenceClass.py


/YOLOv8InferenceClass.py:
--------------------------------------------------------------------------------
  1 | import torch
  2 | import numpy as np
  3 | import cv2
  4 | from time import time
  5 | from ultralytics import YOLO
  6 | 
  7 | import supervision as sv
  8 | 
  9 | 
 10 | class ObjectDetection:
 11 | 
 12 |     def __init__(self, capture_index):
 13 |        
 14 |         self.capture_index = capture_index
 15 |         
 16 |         self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
 17 |         print("Using Device: ", self.device)
 18 |         
 19 |         self.model = self.load_model()
 20 |         
 21 |         self.CLASS_NAMES_DICT = self.model.model.names
 22 |     
 23 |         self.box_annotator = sv.BoxAnnotator(sv.ColorPalette.default(), thickness=3, text_thickness=3, text_scale=1.5)
 24 |     
 25 | 
 26 |     def load_model(self):
 27 |        
 28 |         model = YOLO("yolov8m.pt")  # load a pretrained YOLOv8n model
 29 |         model.fuse()
 30 |     
 31 |         return model
 32 | 
 33 | 
 34 |     def predict(self, frame):
 35 |        
 36 |         results = self.model(frame)
 37 |         
 38 |         return results
 39 |     
 40 | 
 41 |     def plot_bboxes(self, results, frame):
 42 |         
 43 |         xyxys = []
 44 |         confidences = []
 45 |         class_ids = []
 46 |         
 47 |          # Extract detections for person class
 48 |         for result in results:
 49 |             boxes = result.boxes.cpu().numpy()
 50 |             class_id = boxes.cls[0]
 51 |             conf = boxes.conf[0]
 52 |             xyxy = boxes.xyxy[0]
 53 | 
 54 |             if class_id == 0.0:
 55 |           
 56 |               xyxys.append(result.boxes.xyxy.cpu().numpy())
 57 |               confidences.append(result.boxes.conf.cpu().numpy())
 58 |               class_ids.append(result.boxes.cls.cpu().numpy().astype(int))
 59 |             
 60 |         
 61 |         # Setup detections for visualization
 62 |         detections = sv.Detections(
 63 |                     xyxy=results[0].boxes.xyxy.cpu().numpy(),
 64 |                     confidence=results[0].boxes.conf.cpu().numpy(),
 65 |                     class_id=results[0].boxes.cls.cpu().numpy().astype(int),
 66 |                     )
 67 |         
 68 |     
 69 |         # Format custom labels
 70 |         self.labels = [f"{self.CLASS_NAMES_DICT[class_id]} {confidence:0.2f}"
 71 |         for _, confidence, class_id, tracker_id
 72 |         in detections]
 73 |         
 74 |         # Annotate and display frame
 75 |         frame = self.box_annotator.annotate(scene=frame, detections=detections, labels=self.labels)
 76 |         
 77 |         return frame
 78 |     
 79 |     
 80 |     
 81 |     def __call__(self):
 82 | 
 83 |         cap = cv2.VideoCapture(self.capture_index)
 84 |         assert cap.isOpened()
 85 |         cap.set(cv2.CAP_PROP_FRAME_WIDTH, 1280)
 86 |         cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 720)
 87 |       
 88 |         while True:
 89 |           
 90 |             start_time = time()
 91 |             
 92 |             ret, frame = cap.read()
 93 |             assert ret
 94 |             
 95 |             results = self.predict(frame)
 96 |             frame = self.plot_bboxes(results, frame)
 97 |             
 98 |             end_time = time()
 99 |             fps = 1/np.round(end_time - start_time, 2)
100 |              
101 |             cv2.putText(frame, f'FPS: {int(fps)}', (20,70), cv2.FONT_HERSHEY_SIMPLEX, 1.5, (0,255,0), 2)
102 |             
103 |             cv2.imshow('YOLOv8 Detection', frame)
104 |  
105 |             if cv2.waitKey(5) & 0xFF == 27:
106 |                 
107 |                 break
108 |         
109 |         cap.release()
110 |         cv2.destroyAllWindows()
111 |         
112 |         
113 |     
114 | detector = ObjectDetection(capture_index=0)
115 | detector()
116 | 


--------------------------------------------------------------------------------