Glas42 commited on
Commit
0448129
1 Parent(s): a177bf3

Upload ExampleUsage.py

Browse files
Files changed (1) hide show
  1. files/ExampleUsage.py +66 -0
files/ExampleUsage.py ADDED
@@ -0,0 +1,66 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import bettercam
2
+ import pathlib
3
+ import torch
4
+ import time
5
+ import cv2
6
+ import mss
7
+ import os
8
+
9
+ temp = pathlib.PosixPath
10
+ pathlib.PosixPath = pathlib.WindowsPath
11
+
12
+ current_path = os.path.dirname(os.path.abspath(__file__))
13
+ model_path = None
14
+ for file in os.listdir(current_path):
15
+ if file.endswith(".pt"):
16
+ model_path = os.path.join(current_path, file)
17
+ break
18
+ if model_path is None:
19
+ print("No model found. Place it in the same folder as this script.")
20
+ exit()
21
+ model = torch.hub.load('ultralytics/yolov5', 'custom', path=model_path, force_reload=True)
22
+
23
+ sct = mss.mss()
24
+ monitor = sct.monitors[1]
25
+ screen_x = monitor["left"]
26
+ screen_y = monitor["top"]
27
+ screen_width = monitor["width"]
28
+ screen_height = monitor["height"]
29
+
30
+ camera = bettercam.create(output_color="RGB", output_idx=0)
31
+
32
+ cv2.namedWindow('YOLOv5 Detection', cv2.WINDOW_NORMAL)
33
+ cv2.resizeWindow('YOLOv5 Detection', 960, 540)
34
+ cv2.setWindowProperty('YOLOv5 Detection', cv2.WND_PROP_TOPMOST, 1)
35
+
36
+ while True:
37
+ start_time = time.time()
38
+ frame = camera.grab(region=(screen_x, screen_y, screen_x + screen_width, screen_y + screen_height))
39
+ if frame is None: continue
40
+ rgb_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
41
+
42
+ results = model(rgb_frame)
43
+
44
+ boxes = results.pandas().xyxy[0]
45
+
46
+ for _, box in boxes.iterrows():
47
+ label = box['name']
48
+ score = box['confidence']
49
+ x, y, w, h = int(box['xmin']), int(box['ymin']), int(box['xmax'] - box['xmin']), int(box['ymax'] - box['ymin'])
50
+
51
+ if label in ['map']:
52
+ cv2.rectangle(rgb_frame, (x, y), (x + w, y + h), (0, 255, 255), 3)
53
+ cv2.putText(rgb_frame, f"{score:.2f}", (x, y - 10), cv2.FONT_HERSHEY_SIMPLEX, 1.5, (0, 255, 255), 2, cv2.LINE_AA)
54
+ if label in ['arrow']:
55
+ cv2.rectangle(rgb_frame, (x, y), (x + w, y + h), (255, 0, 0), 3)
56
+ cv2.putText(rgb_frame, f"{score:.2f}", (x, y - 10), cv2.FONT_HERSHEY_SIMPLEX, 1.5, (255, 0, 0), 2, cv2.LINE_AA)
57
+
58
+ fps = round(1 / (time.time() - start_time), 1)
59
+ cv2.putText(rgb_frame, f"FPS: {fps}", (20, 60), cv2.FONT_HERSHEY_SIMPLEX, 0.8, (255, 255, 255), 2, cv2.LINE_AA)
60
+ cv2.imshow('YOLOv5 Detection', rgb_frame)
61
+ cv2.resizeWindow('YOLOv5 Detection', 854, 480)
62
+
63
+ if cv2.waitKey(1) & 0xFF == ord('q'):
64
+ break
65
+
66
+ cv2.destroyAllWindows()