Spaces:
Sleeping
Sleeping
| # -*- coding: utf-8 -*- | |
| # @Time : 2022/6/12 15:21 | |
| # @Author : Yajing Zheng | |
| # @File : visualize.py | |
| import cv2 | |
| import numpy as np | |
| import matplotlib | |
| matplotlib.use('Agg') | |
| import matplotlib.pyplot as plt | |
| import matplotlib.cm as cm | |
| import matplotlib.animation as animation | |
| def obtain_spike_video(spikes, video_filename, **dataDict): | |
| spike_h = dataDict.get('spike_h') | |
| spike_w = dataDict.get('spike_w') | |
| timestamps = spikes.shape[0] | |
| mov = cv2.VideoWriter(video_filename, cv2.VideoWriter_fourcc(*'MJPG'), 30, (spike_w, spike_h)) | |
| for iSpk in range(timestamps): | |
| tmpSpk = spikes[iSpk, :, :] * 255 | |
| tmpSpk = cv2.cvtColor(tmpSpk.astype(np.uint8), cv2.COLOR_GRAY2BGR) | |
| mov.write(tmpSpk) | |
| mov.release() | |
| def obtain_reconstruction_video(images, video_filename, **dataDict): | |
| spike_h = dataDict.get('spike_h') | |
| spike_w = dataDict.get('spike_w') | |
| img_num = images.shape[0] | |
| mov = cv2.VideoWriter(video_filename, cv2.VideoWriter_fourcc(*'MJPG'), 30, (spike_w, spike_h)) | |
| for iImg in range(img_num): | |
| tmp_img = images[iImg, :, :] | |
| tmp_img = cv2.cvtColor(tmp_img, cv2.COLOR_GRAY2BGR) | |
| mov.write(tmp_img) | |
| mov.release() | |
| def obtain_mot_video(spikes, video_filename, res_filepath, **dataDict): | |
| spike_h = dataDict.get('spike_h') | |
| spike_w = dataDict.get('spike_w') | |
| gt_file = dataDict.get('labeled_data_dir') | |
| gt_boxes = {} | |
| if gt_file is not None: | |
| gt_f = open(gt_file, 'r') | |
| gt_lines = gt_f.readlines() | |
| for line in gt_lines: | |
| gt_term = line.split(',') | |
| time_step = gt_term[0] | |
| box_id = gt_term[1] | |
| x = float(gt_term[2]) | |
| y = float(gt_term[3]) | |
| w = float(gt_term[4]) | |
| h = float(gt_term[5]) | |
| if str(time_step) not in gt_boxes: | |
| gt_boxes[str(time_step)] = [] | |
| bbox = [box_id, x, y, w, h] | |
| gt_boxes[str(time_step)].append(bbox) | |
| gt_f.close() | |
| result_file = res_filepath | |
| test_boxes = {} | |
| result_f = open(result_file, 'r') | |
| result_lines = result_f.readlines() | |
| color_dict = {} | |
| for line in result_lines: | |
| res_box = line.split(',') | |
| time_step = res_box[0] | |
| track_id = res_box[1] | |
| if track_id not in color_dict.keys(): | |
| colors = (np.random.rand(1, 3) * 255).astype(np.uint8) | |
| color_dict[track_id] = np.squeeze(colors) | |
| x = float(res_box[2]) | |
| y = float(res_box[3]) | |
| w = float(res_box[4]) | |
| h = float(res_box[5]) | |
| if str(time_step) not in test_boxes: | |
| test_boxes[str(time_step)] = [] | |
| test_box = [track_id, x, y, w, h] | |
| test_boxes[str(time_step)].append(test_box) | |
| result_f.close() | |
| mov = cv2.VideoWriter(video_filename, cv2.VideoWriter_fourcc(*'MJPG'), 30, (spike_w, spike_h)) | |
| timestamps = spikes.shape[0] | |
| for t in range(151, timestamps): | |
| # for t in range(160, 1000): | |
| tmp_ivs = spikes[t, :, :] * 255 | |
| tmp_ivs = cv2.cvtColor(tmp_ivs.astype(np.uint8), cv2.COLOR_GRAY2BGR) | |
| if len(gt_boxes) > 0: | |
| if str(t) in gt_boxes: | |
| gts = gt_boxes[str(t)] | |
| gt_num = len(gts) | |
| for i in range(gt_num): | |
| box = gts[i] | |
| box_id = box[0] | |
| cv2.rectangle(tmp_ivs, (int(box[2]), int(box[1])), | |
| (int(box[2] + box[4]), int(box[1] + box[3])), | |
| (int(255), int(255), int(255)), 2) | |
| if str(t) in test_boxes: | |
| test = test_boxes[str(t)] | |
| test_num = len(test) | |
| for i in range(test_num): | |
| box = test[i] | |
| box_id = box[0] | |
| colors = color_dict[box_id] | |
| cv2.rectangle(tmp_ivs, (int(box[2]), int(box[1])), | |
| (int(box[2] + box[4]), int(box[1] + box[3])), | |
| (int(colors[0]), int(colors[1]), int(colors[2])), 2) | |
| mov.write(tmp_ivs) | |
| mov.release() | |
| def obtain_detection_video(spikes, video_filename, res_filepath, evaluate_seq_len, begin_idx=0, **dataDict): | |
| spike_h = dataDict.get('spike_h') | |
| spike_w = dataDict.get('spike_w') | |
| gt_file = dataDict.get('labeled_data_dir') | |
| gt_boxes = {} | |
| if gt_file is not None: | |
| start_idx = begin_idx | |
| end_idx = begin_idx + evaluate_seq_len | |
| for seq_no in range(start_idx, end_idx): | |
| gt_filename = gt_file[seq_no] | |
| gt_f = open(gt_filename, 'r') | |
| gt_lines = gt_f.readlines() | |
| for line in gt_lines: | |
| tmp_box = line.split(',') | |
| x = float(tmp_box[0]) | |
| y = float(tmp_box[1]) | |
| w = float(tmp_box[2]) | |
| h = float(tmp_box[3]) | |
| box_id = int(0) | |
| if str(seq_no) not in gt_boxes: | |
| gt_boxes[str(seq_no)] = [] | |
| bbox = [box_id, x, y, w, h] | |
| gt_boxes[str(seq_no)].append(bbox) | |
| gt_f.close() | |
| result_file = res_filepath | |
| test_boxes = {} | |
| result_f = open(result_file, 'r') | |
| result_lines = result_f.readlines() | |
| color_dict = {} | |
| for line in result_lines: | |
| res_box = line.split(',') | |
| time_step = res_box[0] | |
| track_id = res_box[1] | |
| if track_id not in color_dict.keys(): | |
| colors = (np.random.rand(1, 3) * 255).astype(np.uint8) | |
| color_dict[track_id] = np.squeeze(colors) | |
| x = float(res_box[2]) | |
| y = float(res_box[3]) | |
| w = float(res_box[4]) | |
| h = float(res_box[5]) | |
| if str(time_step) not in test_boxes: | |
| test_boxes[str(time_step)] = [] | |
| test_box = [track_id, x, y, w, h] | |
| test_boxes[str(time_step)].append(test_box) | |
| result_f.close() | |
| mov = cv2.VideoWriter(video_filename, cv2.VideoWriter_fourcc(*'MJPG'), 30, (spike_w, spike_h)) | |
| block_len = spikes.shape[0] | |
| # gt_intv = int(block_len/evaluate_seq_len) | |
| gt_intv = 400 | |
| # for t in range(150, block_len): | |
| for i_gt in range(start_idx+1, end_idx): | |
| t = i_gt * gt_intv + int(gt_intv/2) | |
| tmp_ivs = spikes[t, :, :] * 255 | |
| tmp_ivs = cv2.cvtColor(tmp_ivs.astype(np.uint8), cv2.COLOR_GRAY2BGR) | |
| if len(gt_boxes) > 0: | |
| gts = gt_boxes[str(i_gt)] | |
| gt_num = len(gts) | |
| for i in range(gt_num): | |
| box = gts[i] | |
| cv2.rectangle(tmp_ivs, (int(spike_w - box[1]), int(box[2])), | |
| (int(spike_w - box[1] - box[3]), int(box[2] + box[4])), | |
| (int(255), int(255), int(255)), 2) | |
| if str(t) in test_boxes: | |
| test = test_boxes[str(t)] | |
| test_num = len(test) | |
| for i in range(test_num): | |
| box = test[i] | |
| box_id = box[0] | |
| colors = color_dict[box_id] | |
| cv2.rectangle(tmp_ivs, (int(box[2]), int(box[1])), | |
| (int(box[2] + box[4]), int(box[1] + box[3])), | |
| (int(colors[0]), int(colors[1]), int(colors[2])), 2) | |
| mov.write(tmp_ivs) | |
| mov.release() | |
| def get_heatVideo(results, video_filename): | |
| results = np.array(results) | |
| frame_num = results.shape[0] | |
| frames = [] | |
| fig = plt.figure() | |
| for i in range(frame_num): | |
| tmp_res = results[i] | |
| # frames.append([plt.imshow(tmp_res, cmap=cm.Greys_r, animated=True)]) | |
| frames.append([plt.imshow(tmp_res, cmap=cm.Blues, animated=True)]) | |
| ani = animation.ArtistAnimation(fig, frames, interval=50, blit=True, | |
| repeat_delay=1000) | |
| # change the path to where you save ffmpeg | |
| plt.rcParams['animation.ffmpeg_path'] = 'F:\\ffmpeg-N-99818-g993429cfb4-win64-gpl-shared-vulkan\\bin\\ffmpeg.exe' | |
| FFwrite = animation.FFMpegWriter(fps=30, extra_args=['-vcodec', 'libx264']) | |
| ani.save(video_filename, writer=FFwrite) | |
| plt.show() | |