This repository was archived by the owner on May 27, 2024. It is now read-only.
-
Notifications
You must be signed in to change notification settings - Fork 2
Expand file tree
/
Copy pathworker.py
More file actions
96 lines (73 loc) · 3.04 KB
/
worker.py
File metadata and controls
96 lines (73 loc) · 3.04 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
import glob
import os
import time
from datetime import datetime
import cv2
import numpy as np
from PySide6.QtCore import Signal, QThread, QDir
from ultralytics import YOLO
from config import allow_extensions
class Worker(QThread):
update_progress = Signal(int)
end_working = Signal(list)
progress_increment = Signal()
def __init__(self, path_to_file: str, save_result: bool):
self.path_to_file = path_to_file
self.save_result = save_result
super().__init__()
@staticmethod
def is_image_file(filename):
return any(extension in filename for extension in allow_extensions)
def find_paths(self, path):
result = []
if os.path.isfile(path):
if self.is_image_file(path):
result.append(path)
elif os.path.isdir(path):
result = []
for ext in allow_extensions:
result += glob.glob(os.path.join(path, f'**/*{ext}'), recursive=True)
return result
def run(self):
paths = self.find_paths(self.path_to_file)
if not QDir().exists("results"):
QDir().mkdir("results")
if len(paths) == 0:
self.update_progress.emit(100)
self.end_working.emit(None)
return
else:
for i in range(51):
time.sleep(0.0015)
self.update_progress.emit(i)
model = YOLO('./weights/best.pt')
results = model(paths)
response = []
for result in results:
path = result.path
filename = os.path.basename(path)
for box in result[0].boxes:
class_id = result.names[box.cls[0].item()]
cords = box.xyxy[0].tolist()
cords = [round(x) for x in cords]
img = cv2.imread(path)
img = np.array(img) # Переводим в numpy массивы
img = img.astype('float32') # Изменяем на тип float
img = img / 255.0 # Нормализируем пиксили
cv2.rectangle(img, (cords[0], cords[1]), (cords[2], cords[3]), (0, 255, 255), 5)
if self.save_result:
foldername = str(datetime.now())[:-7].replace(":", "-")
QDir().mkdir(f"./results/{foldername}")
cv2.imwrite(f'./results/{foldername}/{filename}', img)
if class_id != "не дефект":
response.append({
"type": class_id,
"image": img,
"filename": filename,
"cors": f"X0: {cords[0]}, Y0: {cords[1]}; X1: {cords[2]}, Y1: {cords[3]}"
})
for i in range(50//len(results)):
self.progress_increment.emit()
time.sleep(0.015)
self.update_progress.emit(100)
self.end_working.emit(response)