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)