lyxxxh2
71 天前
```
# 以下由 gpt 生成
def init_input_and_output_door(self):
frame_results = copy.deepcopy(
self.frame_results
) # self.frame_results 会被以下代码改变 找太麻烦了 直接拷贝到新对象
door_line = [[828, 438], [1208, 868]]
door_area_line = [[623, 523], [1144, 982]]
door_line = [[336, 0], [1698, 1222]]
door_line =[[186, 0], [1428, 1222]]
door_line = [[336, 0], [1698, 1222]]
def is_above_line(point, line):
(x1, y1), (x2, y2) = line
return (y2 - y1) * (point[0] - x1) > (x2 - x1) * (point[1] - y1)
def filter_boxes_above_line(frame_results, door_area_line):
filtered_results = []
for frame in frame_results:
filtered_boxes = []
for box in frame["boxes"]:
x_min, y_min, x_max, y_max = box
box_center = [(x_min + x_max) / 2, (y_min + y_max) / 2]
if is_above_line(box_center, door_area_line):
filtered_boxes.append(box)
if filtered_boxes:
filtered_frame = {
"id": frame["id"],
"boxes": filtered_boxes,
"reid_dets": frame["reid_dets"], # 保留每个帧的 reid_dets
}
filtered_results.append(filtered_frame)
return filtered_results
filtered_frame_results = filter_boxes_above_line(frame_results, door_area_line)
filtered_frame_results = [
item for item in filtered_frame_results if len(item["boxes"]) != 0
]
# 使用 IOUTracker 类
tracker = IOUTracker(iou_threshold=0.3)
all_tracks = tracker.track_objects(filtered_frame_results)
# 用于记录人的信息、reid_dets 和 frame_ids
person_groups = defaultdict(list)
person_reid_dets = defaultdict(list)
person_frame_ids = defaultdict(list)
for frame, track_ids in zip(filtered_frame_results, all_tracks):
for box in frame["boxes"]:
matching_id = [
id for id, tracked_box in track_ids.items() if tracked_box == box
]
if matching_id:
box.append("person_id: {}".format(matching_id[0]))
person_groups[matching_id[0]].append(box[:-1]) # 不包括 person_id
person_reid_dets[matching_id[0]].append(
frame["reid_dets"]
) # 记录 reid_dets
person_frame_ids[matching_id[0]].append(
frame["id"]
) # 记录 frame id
crossing_events = {}
final_crossing_frames = {}
# 门线参数
m = (door_line[1][1] - door_line[0][1]) / (
door_line[1][0] - door_line[0][0]
) # 斜率
c = door_line[0][1] - m * door_line[0][0] # 截距
def check_position(x, y):
line_y = m * x + c
return y - line_y # > 0 在线上方，< 0 在线下方
for person_id, boxes in person_groups.items():
last_pos = None
entry_index = None
exit_index = None
# if person_id != 2:
# continue
for idx, box in enumerate(boxes):
x_center = (box[0] + box[2]) / 2
y_center = (box[1] + box[3]) / 2
current_pos = check_position(x_center, y_center)
if last_pos is not None:
if last_pos > 0 and current_pos < 0:
exit_index = idx
elif last_pos < 0 and current_pos > 0:
entry_index = idx
last_pos = current_pos
crossing_events[person_id] = {
"input_index": exit_index,
"output_index": entry_index,
}
for person_id, events in crossing_events.items():
person_data = []
boxes = person_groups[person_id]
reid_dets = person_reid_dets[person_id]
frame_ids = person_frame_ids[person_id] # 获取 frame id 列表
if events["input_index"] is not None:
input_index = events["input_index"]
input_frames = boxes[input_index : input_index + 3]
input_reid_dets = reid_dets[input_index : input_index + 3]
input_frame_ids = frame_ids[
input_index : input_index + 3
] # 取相应的 frame ids
person_data.append(
{
"type": "output",
"all_box": input_frames,
"reid_dets": input_reid_dets,
"frame_ids": input_frame_ids,
"ids": list(
range(input_index, input_index + len(input_frames))
),
}
)
if events["output_index"] is not None:
output_index = events["output_index"]
output_frames_start = max(0, output_index - 2)
output_frames = boxes[output_frames_start : output_index + 1]
output_reid_dets = reid_dets[output_frames_start : output_index + 1]
output_frame_ids = frame_ids[
output_frames_start : output_index + 1
] # 取相应的 frame ids
person_data.append(
{
"type": "input",
"all_box": output_frames,
"reid_dets": output_reid_dets,
"frame_ids": output_frame_ids,
"ids": list(range(output_frames_start, output_index + 1)),
}
)
final_crossing_frames[person_id] = person_data
```
之前让 gpt 写的。
也不是不能用,就是没一点"设计"的感觉。
让他再优化下,直接处 bug 。
逻辑越复杂,我就不敢用 ai 。
因为描述太麻烦了,得写篇作文给他。
再者很难阅读和修改。
后面我直接重写了
```
# 定义数据
for box, person_id in zip(data["head_boxes"], person_ids):
if person_id is None:
continue
if person_id not in self.door: # 初始化 key
self.door[person_id] = []
self.user_door_status[person_id] = None # 用户和门状态
for device in self.devices: # 用户和设备状态
if person_id not in device["persons"]:
device["persons"][person_id] = []
self.door[person_id].append(
{"box": box, "up_or_line": self.box_line_up_or_down(box)}
) # 添加到门记录器
if len(self.door[person_id]) > 100: # 用户超过一百
self.door[person_id] = self.door[person_id][-50:]
for device in self.devices: # 添加到设备记录器
if person_id not in device:
device["persons"][person_id].append(
self.calculate_containment_ratio(device["box"], box)
)
if len(device["persons"][person_id]) > 100:
device["persons"][person_id] = device["persons"][person_id][-50:]
# 根据
for person_id in self.door:
...
....
```