def show_current_frame(self,
normboxes,
class_names,
window_name='current'):
Memory.current_frame_rails = []
# normboxes = normboxes[::-1]
top_list, bottom_list = self.video_loader.get_top_bottom_list()
# top_list, bottom_list = top_list[::-1], bottom_list[::-1]
output_image = self.video_loader.get_image()
no_boxes = True
timer1 = time.time()
# normboxes = self.init_current_rails(normboxes, output_image)
# self.init_current_rails(normboxes, output_image)
if not Memory.current_rails_detected:
self.init_current_rails(normboxes, "central rails", output_image)
else:
self.continue_current_rails(normboxes, output_image)
normboxes = Memory.current_frame_rails
# print(time.time() - timer1)
for box in normboxes:
###############
# вывод
###############
box[4] = "current rails"
if box[4] == "current rails":
display_box(output_image, [box[0], box[1], box[2], box[3]],
color=(0, 255, 0),
thickness=3)
text = box[4]
cv2.putText(output_image,
text, (int(box[0]), int(box[1]) + 10),
self.font,
0.5, (0, 255, 0),
thickness=2)
else:
box_class = [
"double central rails", "double left rails",
"double right rails", "central rails", "left rails",
"right rails", "half left rails", "half right rails",
"switch right back", "switch left back",
"switch right front", "switch left front", "switch",
"left crossing", "right crossing", "double cross"
]
display_box(output_image, [box[0], box[1], box[2], box[3]],
color=self.class_colours[box_class.index(box[4])],
thickness=3)
text = box[4]
cv2.putText(output_image,
text, (int(box[0]), int(box[1]) + 10),
self.font,
0.5,
self.class_colours[box_class.index(box[4])],
thickness=2)
return output_image
def show_current_frame(self, norm_boxes, class_names, window_name='current'):
output_image = self.video_loader.get_image()
frame_width = InputPlug.FrameWidth
frame_height = InputPlug.FrameHeight
rail_boxes = norm_boxes
crop_coors = np.int16(InputPlug.CropCoors * frame_height)
timer1 = time.time()
filtred_boxes = analytics(rail_boxes, crop_coors, frame_width, frame_height, output_image)
# polygon = analytics(rail_boxes, crop_coors, frame_width, frame_height, output_image)
# print(time.time() - timer1)
no_boxes = True
# uncomment
# вывод боксов
for box in norm_boxes:
box_class = int(box[4])
class_name = class_names[box_class]
###############
# вывод
###############
display_box(output_image, [box[0], box[1], box[2], box[3]],
color=self.class_colours[box_class], thickness=1)
text = class_names[box_class]
cv2.putText(output_image, text, (int(box[0]), int(box[1]) + 10),
self.font, 0.5, self.class_colours[box_class], thickness=1)
#
for box in filtred_boxes:
display_box(output_image, [box[0], box[1],
box[2], box[3]], color=(0, 255, 0), thickness=2)
return output_image
def display_line(self, line, boxes, class_names, height_line=64):
# boxes = boxes[::-1]
line = cv2.resize(line, dsize=(line.shape[1], height_line), interpolation=cv2.INTER_CUBIC)
for box in boxes:
box_class = int(box[5])
temp = box[0:4].to(dtype=torch.int32).tolist()
display_box(line, [temp[0], 0, temp[2], height_line - 1], color=self.class_colours[box_class], thickness=30)
cv2.putText(line, class_names[box_class], (temp[0], 10),
self.font, 0.5, self.class_colours[box_class], thickness=2)
return line
def make_data(self, boxes, frame_name, class_names, window_name='display'):
top_list, bottom_list = self.video_loader.get_top_bottom_list()
output_image = self.video_loader.get_image()
no_boxes = True
df = pd.DataFrame(columns=['label', 'cx', 'cy'])
# pred_lists = []
indexes = 0
for i in range(0, len(boxes)):
if len(boxes[i]) != 0:
no_boxes = False
y1 = top_list[i]
y2 = bottom_list[i]
for box in boxes[i]:
box_class = int(box[5])
temp = box[0:3].to(dtype=torch.int32)
cx = temp[0].item() + int(
(temp[2].item() - temp[0].item()) / 2)
cy = y1 + int((y2 - y1) / 2)
# pred_list = [class_names[box_class], cx, cy, indexes]
# pred_lists.append(pred_list)
pred_list = [class_names[box_class], cx, cy]
df.loc[indexes] = pred_list
display_box(output_image, [temp[0], y1, temp[2], y2],
color=self.class_colours[box_class],
thickness=3)
text = class_names[box_class]
cv2.putText(output_image,
text, (temp[0], y1 + 10),
self.font,
0.5,
self.class_colours[box_class],
thickness=2)
cv2.putText(output_image,
str(indexes), (temp[2], y1 + 10),
self.font,
0.5, (0, 0, 0),
thickness=1)
indexes += 1
cv2.imwrite(
'/media/user3070/data/teplovoz_table_data/video2/' + frame_name +
'.png', output_image)
df.to_csv('/media/user3070/data/teplovoz_table_data/video2/' +
frame_name + '.csv',
index_label='num')
cv2.imshow(window_name, output_image)
cv2.waitKey(0)
def make_norm_data(self,
norm_boxes,
frame_name,
class_names,
window_name='display'):
output_image = self.video_loader.get_image()
no_boxes = True
df = pd.DataFrame(
norm_boxes, columns=['x0', 'y0', 'x1', 'y1', 'class', 'num_class'])
# pred_lists = []
indexes = 0
for box in norm_boxes:
# pred_list = [class_names[box_class], cx, cy, indexes]
# pred_lists.append(pred_list)
pred_list = box
df.loc[indexes] = pred_list
if box[4] == "current rails":
display_box(output_image, [box[0], box[1], box[2], box[3]],
color=(0, 255, 0),
thickness=3)
text = box[4]
cv2.putText(output_image,
text, (int(box[0]), int(box[1]) + 10),
self.font,
0.5, (0, 255, 0),
thickness=2)
else:
box_class = [
"double central rails", "double left rails",
"double right rails", "central rails", "left rails",
"right rails", "half left rails", "half right rails",
"switch right back", "switch left back",
"switch right front", "switch left front", "switch",
"left crossing", "right crossing", "double cross"
]
display_box(output_image, [box[0], box[1], box[2], box[3]],
color=self.class_colours[box_class.index(box[4])],
thickness=3)
text = box[4]
cv2.putText(output_image,
text, (int(box[0]), int(box[1]) + 10),
self.font,
0.5,
self.class_colours[box_class.index(box[4])],
thickness=2)
cv2.imwrite('video/' + frame_name + '.png', output_image)
df.to_csv('video2/' + frame_name + '.csv', index_label='num')
def display_train(self, calc_bilder, opt_alg):
# color = (255, 0, 0)
tracks = opt_alg["tracks"]
for trk in tracks:
for box in trk.history_box:
display_box(self.im, box, trk.color_box)
centr = [] # calculate_centr_box(trk.history_box, np.array([0.5, 0.5]))
for box in trk.history_box:
centr.append(
(box[2:4] + box[:2]) / 2
)
for i in range(1, len(centr)):
cv2.line(self.im, tuple(np.int32(centr[i - 1])),
tuple(np.int32(centr[i])), trk.color_box, thickness=3)