Skip to content

Commit 6d64e40

Browse files
Koldim2001Koldim2001
committed
pep8 fix code
1 parent 481b530 commit 6d64e40

File tree

10 files changed

+167
-100
lines changed

10 files changed

+167
-100
lines changed

‎main.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -22,9 +22,9 @@ def main(config) -> None:
2222

2323
if send_info_db:
2424
send_info_db_node = SendInfoDBNode(config)
25-
25+
2626
for frame_element in video_reader.process():
27-
27+
2828
frame_element = detection_node.process(frame_element)
2929
frame_element = tracker_info_update_node.process(frame_element)
3030
frame_element = calc_statistics_node.process(frame_element)

‎main_optimized.py

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,9 @@
1-
import hydra
21
from time import sleep, time
3-
from tqdm import tqdm
42
from multiprocessing import Process, Queue
53

4+
import hydra
5+
from tqdm import tqdm
6+
67
from nodes.VideoReader import VideoReader
78
from nodes.ShowNode import ShowNode
89
from nodes.VideoSaverNode import VideoSaverNode
@@ -109,7 +110,7 @@ def main(config) -> None:
109110
target=proc_show_node,
110111
args=(queue_track_update_out, config),
111112
name="proc_show_node",
112-
)
113+
),
113114
]
114115

115116
for p in processes:

‎nodes/CalcStatisticsNode.py

Lines changed: 19 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -5,17 +5,22 @@
55
from elements.VideoEndBreakElement import VideoEndBreakElement
66
from utils_local.utils import profile_time
77

8+
89
class CalcStatisticsNode:
910
"""Модуль для расчета загруженности дорог (вычисление статистик)"""
1011

1112
def __init__(self, config: dict) -> None:
1213
config_general = config["general"]
1314

14-
self.time_buffer_analytics = config_general["buffer_analytics"] # размер времени буфера в минутах
15-
self.min_time_life_track = config_general["min_time_life_track"] # минимальное время жизни трека в сек
15+
self.time_buffer_analytics = config_general[
16+
"buffer_analytics"
17+
] # размер времени буфера в минутах
18+
self.min_time_life_track = config_general[
19+
"min_time_life_track"
20+
] # минимальное время жизни трека в сек
1621
self.count_cars_buffer_frames = config_general["count_cars_buffer_frames"]
1722
self.cars_buffer = deque(maxlen=self.count_cars_buffer_frames) # создали буфер значений
18-
23+
1924
@profile_time
2025
def process(self, frame_element: FrameElement) -> FrameElement:
2126
# Выйти из обработки если это пришел VideoEndBreakElement а не FrameElement
@@ -25,19 +30,25 @@ def process(self, frame_element: FrameElement) -> FrameElement:
2530
frame_element, FrameElement
2631
), f"CalcStatisticsNode | Неправильный формат входного элемента {type(frame_element)}"
2732

28-
2933
buffer_tracks = frame_element.buffer_tracks
3034
self.cars_buffer.append(len(frame_element.id_list))
3135

3236
info_dictionary = {}
33-
info_dictionary['cars_amount'] = round(np.mean(self.cars_buffer))
34-
roads_activity = {1:0, 2:0, 3:0, 4:0, 5:0} # всего 5 дорог (занулим стартовое значение)
37+
info_dictionary["cars_amount"] = round(np.mean(self.cars_buffer))
38+
roads_activity = {
39+
1: 0,
40+
2: 0,
41+
3: 0,
42+
4: 0,
43+
5: 0,
44+
} # всего 5 дорог (занулим стартовое значение)
3545

3646
# Посчитаем чило машин которые довно живут и имеют значения дороги приезда
3747
for _, track_element in buffer_tracks.items():
3848
if (
39-
track_element.timestamp_last - track_element.timestamp_init_road > self.min_time_life_track and
40-
track_element.start_road is not None
49+
track_element.timestamp_last - track_element.timestamp_init_road
50+
> self.min_time_life_track
51+
and track_element.start_road is not None
4152
):
4253
key = track_element.start_road
4354
roads_activity[key] += 1
@@ -52,4 +63,3 @@ def process(self, frame_element: FrameElement) -> FrameElement:
5263
frame_element.info = info_dictionary
5364

5465
return frame_element
55-

‎nodes/DetectionTrackingNodes.py

Lines changed: 20 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@
77
from elements.VideoEndBreakElement import VideoEndBreakElement
88
from byte_tracker.byte_tracker_model import BYTETracker as ByteTracker
99

10+
1011
class DetectionTrackingNodes:
1112
"""Модуль инференса модели детекции + трекинг алгоритма"""
1213

@@ -31,7 +32,9 @@ def __init__(self, config) -> None:
3132
match_thresh = config_bytetrack["match_thresh"]
3233
track_buffer = config_bytetrack["track_buffer"]
3334
fps = 30 # ставим равным 30 чтобы track_buffer мерился в кадрах
34-
self.tracker = ByteTracker(fps, first_track_thresh, second_track_thresh, match_thresh, track_buffer, 1)
35+
self.tracker = ByteTracker(
36+
fps, first_track_thresh, second_track_thresh, match_thresh, track_buffer, 1
37+
)
3538

3639
@profile_time
3740
def process(self, frame_element: FrameElement) -> FrameElement:
@@ -82,14 +85,23 @@ def _get_results_dor_tracker(self, results) -> np.ndarray:
8285
class_id = result.boxes.cls.cpu().numpy().astype(int)
8386
# трекаем те же классы что и детектируем
8487
if class_id[0] in self.classes_to_detect:
85-
88+
8689
bbox = result.boxes.xyxy.cpu().numpy()
8790
confidence = result.boxes.conf.cpu().numpy()
8891

89-
class_id_value = 2 # Будем все трекуемые объекты считать классом car чтобы не было ошибок
90-
91-
merged_detection = [bbox[0][0], bbox[0][1], bbox[0][2], bbox[0][3], confidence[0], class_id_value]
92-
92+
class_id_value = (
93+
2 # Будем все трекуемые объекты считать классом car чтобы не было ошибок
94+
)
95+
96+
merged_detection = [
97+
bbox[0][0],
98+
bbox[0][1],
99+
bbox[0][2],
100+
bbox[0][3],
101+
confidence[0],
102+
class_id_value,
103+
]
104+
93105
detections_list.append(merged_detection)
94-
95-
return np.array(detections_list)
106+
107+
return np.array(detections_list)

‎nodes/SendInfoDBNode.py

Lines changed: 31 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@
88

99
logger = logging.getLogger(__name__)
1010

11+
1112
class SendInfoDBNode:
1213
"""Модуль для отправки актуальной информации о трафике в базу данных"""
1314

@@ -81,7 +82,6 @@ def __init__(self, config: dict) -> None:
8182
f"Error while creating table: {error}"
8283
)
8384

84-
8585
@profile_time
8686
def process(self, frame_element: FrameElement) -> FrameElement:
8787
# Выйти из обработки если это пришел VideoEndBreakElement а не FrameElement
@@ -101,11 +101,12 @@ def process(self, frame_element: FrameElement) -> FrameElement:
101101
if current_time - self.last_db_update >= self.how_often_add_info:
102102
self._insert_in_db(info_dictionary, timestamp, timestamp_date)
103103
frame_element.send_info_of_frame_to_db = True
104-
self.last_db_update = current_time # Обновление времени последнего обновления базы данных
104+
self.last_db_update = (
105+
current_time # Обновление времени последнего обновления базы данных
106+
)
105107

106108
return frame_element
107109

108-
109110
def _insert_in_db(self, info_dictionary: dict, timestamp: float, timestamp_date: float) -> None:
110111
# Формирование и выполнение SQL-запроса для вставки данных в бд
111112
insert_query = (
@@ -119,13 +120,33 @@ def _insert_in_db(self, info_dictionary: dict, timestamp: float, timestamp_date:
119120
(
120121
timestamp,
121122
timestamp_date,
122-
info_dictionary['cars_amount'],
123-
info_dictionary['roads_activity'][1] if timestamp >= self.buffer_analytics_sec else None,
124-
info_dictionary['roads_activity'][2] if timestamp >= self.buffer_analytics_sec else None,
125-
info_dictionary['roads_activity'][3] if timestamp >= self.buffer_analytics_sec else None,
126-
info_dictionary['roads_activity'][4] if timestamp >= self.buffer_analytics_sec else None,
127-
info_dictionary['roads_activity'][5] if timestamp >= self.buffer_analytics_sec else None,
128-
)
123+
info_dictionary["cars_amount"],
124+
(
125+
info_dictionary["roads_activity"][1]
126+
if timestamp >= self.buffer_analytics_sec
127+
else None
128+
),
129+
(
130+
info_dictionary["roads_activity"][2]
131+
if timestamp >= self.buffer_analytics_sec
132+
else None
133+
),
134+
(
135+
info_dictionary["roads_activity"][3]
136+
if timestamp >= self.buffer_analytics_sec
137+
else None
138+
),
139+
(
140+
info_dictionary["roads_activity"][4]
141+
if timestamp >= self.buffer_analytics_sec
142+
else None
143+
),
144+
(
145+
info_dictionary["roads_activity"][5]
146+
if timestamp >= self.buffer_analytics_sec
147+
else None
148+
),
149+
),
129150
)
130151
self.connection.commit()
131152
logger.info(

0 commit comments

Comments
 (0)