nehulagrawal commited on
Commit
88eaa86
·
1 Parent(s): 544b98c

Delete detection.py

Browse files
Files changed (1) hide show
  1. detection.py +0 -59
detection.py DELETED
@@ -1,59 +0,0 @@
1
- import cv2
2
- import IPython
3
- from PIL import ImageColor
4
- from ultralytics import YOLO
5
-
6
- class ObjectDetection:
7
- def __init__(self, model_name='Yolov8'):
8
- self.model_name = model_name
9
- self.model = self.load_model()
10
- self.classes = self.model.names
11
- self.device = 'cpu'
12
-
13
- def load_model(self):
14
- model = YOLO(f"weights/{self.model_name}_best.pt")
15
- return model
16
-
17
- def v8_score_frame(self, frame):
18
- results = self.model(frame)
19
-
20
- labels = []
21
- confidences = []
22
- coords = []
23
-
24
- for result in results:
25
- boxes = result.boxes.cpu().numpy()
26
-
27
- label = boxes.cls
28
- conf = boxes.conf
29
- coord = boxes.xyxy
30
-
31
- labels.extend(label)
32
- confidences.extend(conf)
33
- coords.extend(coord)
34
-
35
- return labels, confidences, coords
36
-
37
- def get_coords(self, frame, row):
38
- return int(row[0]), int(row[1]), int(row[2]), int(row[3])
39
-
40
- def class_to_label(self, x):
41
- return self.classes[int(x)]
42
-
43
- def get_color(self, code):
44
- rgb = ImageColor.getcolor(code, "RGB")
45
- return rgb
46
-
47
- def plot_bboxes(self, results, frame, threshold=0.5, box_color='orange', text_color='black'):
48
- labels, conf, coord = results
49
- box_color = self.get_color(box_color)
50
- text_color = self.get_color(text_color)
51
-
52
- for i in range(len(labels)):
53
- if conf[i] >= threshold:
54
- x1, y1, x2, y2 = self.get_coords(frame, coord[i])
55
- class_name = self.class_to_label(labels[i])
56
-
57
- cv2.rectangle(frame, (x1, y1), (x2, y2), box_color, 2)
58
- cv2.putText(frame, f"{class_name} - {conf[i]*100:.2f}%", (x1, y1), cv2.FONT_HERSHEY_COMPLEX, 0.5, text_color)
59
- return frame