企业官网建设流程全解析

目录

极小目标检测：

csv进行可视化；

极小目标检测：

TrackNetV3

推理代码

# # Predict the label csv from the video (后面的--large_video是可选的) # python predict.py --video_file test.mp4 --tracknet_file ckpts/TrackNet_best.pt --inpaintnet_file ckpts/InpaintNet_best.pt --save_dir prediction --large_video # # Predict the label csv from the video, and output a video with predicted trajectory (后面的--large_video是可选的) # python predict.py --video_file test.mp4 --tracknet_file ckpts/TrackNet_best.pt --inpaintnet_file ckpts/InpaintNet_best.pt --save_dir prediction --output_video --large_video # # For large video # # Enable the flag to use an IterableDataset instead of the normal Dataset, which prevents memory errors. Note that this will decrease the inference speed.--large_video # # Use to set the number of samples for background estimation.--max_sample_num # # Use to specify the start and end seconds of the video for background estimation.--video_range # python predict.py --video_file test.mp4 --tracknet_file ckpts/TrackNet_best.pt --inpaintnet_file ckpts/InpaintNet_best.pt --save_dir prediction --large_video --video_range 324,330 import os import argparse import numpy as np from tqdm import tqdm import torch from torch.utils.data import DataLoader from test import predict_location, get_ensemble_weight, generate_inpaint_mask from dataset import Shuttlecock_Trajectory_Dataset, Video_IterableDataset from utils.general import * def predict(indices, y_pred=None, c_pred=None, img_scaler=(1, 1)): """ Predict coordinates from heatmap or inpainted coordinates. Args: indices (torch.Tensor): indices of input sequence with shape (N, L, 2) y_pred (torch.Tensor, optional): predicted heatmap sequence with shape (N, L, H, W) c_pred (torch.Tensor, optional): predicted inpainted coordinates sequence with shape (N, L, 2) img_scaler (Tuple): image scaler (w_scaler, h_scaler) Returns: pred_dict (Dict): dictionary of predicted coordinates Format: {'Frame':[], 'X':[], 'Y':[], 'Visibility':[]} """ pred_dict = {'Frame':[], 'X':[], 'Y':[], 'Visibility':[]} batch_size, seq_len = indices.shape[0], indices.shape[1] indices = indices.detach().cpu().numpy()if torch.is_tensor(indices) else indices.numpy() # Transform input for heatmap prediction if y_pred is not None: y_pred = y_pred > 0.5 y_pred = y_pred.detach().cpu().numpy() if torch.is_tensor(y_pred) else y_pred y_pred = to_img_format(y_pred) # (N, L, H, W) # Transform input for coordinate prediction if c_pred is not None: c_pred = c_pred.detach().cpu().numpy() if torch.is_tensor(c_pred) else c_pred prev_f_i = -1 for n in range(batch_size): for f in range(seq_len): f_i = indices[n][f][1] if f_i != prev_f_i: if c_pred is not None: # Predict from coordinate c_p = c_pred[n][f] cx_pred, cy_pred = int(c_p[0] * WIDTH * img_scaler[0]), int(c_p[1] * HEIGHT* img_scaler[1]) elif y_pred is not None: # Predict from heatmap y_p = y_pred[n][f] bbox_pred = predict_location(to_img(y_p)) cx_pred, cy_pred = int(bbox_pred[0]+bbox_pred[2]/2), int(bbox_pred[1]+bbox_pred[3]/2) cx_pred, cy_pred = int(cx_pred*img_scaler[0]), int(cy_pred*img_scaler[1]) else: raise ValueError('Invalid input') vis_pred = 0 if cx_pred == 0 and cy_pred == 0 else 1 pred_dict['Frame'].append(int(f_i)) pred_dict['X'].append(cx_pred) pred_dict['Y'].append(cy_pred) pred_dict['Visibility'].append(vis_pred) prev_f_i = f_i else: break return pred_dict if __name__ == '__main__': parser = argparse.ArgumentParser() parser.add_argument('--video_file', type=str, default=r"C:\Users\ChanJing-01\Videos\yumao\yumao1.mp4", help='file path of the video') parser.add_argument('--tracknet_file', type=str, default=r"D:\project_2026\yumaoqiu\TrackNetV3-master\ckpts\TrackNet_best.pt", help='file path of the TrackNet model checkpoint') parser.add_argument('--inpaintnet_file', type=str, default=r'D:\project_2026\yumaoqiu\TrackNetV3-master\ckpts\InpaintNet_best.pt', help='file path of the InpaintNet model checkpoint') # parser.add_argument('--video_file', type=str, help='file path of the video') # parser.add_argument('--tracknet_file', type=str, help='file path of the TrackNet model checkpoint') # parser.add_argument('--inpaintnet_file', type=str, default='', help='file path of the InpaintNet model checkpoint') parser.add_argument('--batch_size', type=int, default=16, help='batch size for inference') parser.add_argument('--eval_mode', type=str, default='weight', choices=['nonoverlap', 'average', 'weight'], help='evaluation mode') parser.add_argument('--max_sample_num', type=int, default=1800, help='maximum number of frames to sample for generating median image') parser.add_argument('--video_range', type=lambda splits: [int(s) for s in splits.split(',')], default=None, help='range of start second and end second of the video for generating median image') parser.add_argument('--save_dir', type=str, default='pred_result', help='directory to save the prediction result') parser.add_argument('--large_video', action='store_true', default=True, help='whether to process large video') parser.add_argument('--output_video', action='store_true', default=True, help='whether to output video with predicted trajectory') parser.add_argument('--traj_len', type=int, default=8, help='length of trajectory to draw on video') args = parser.parse_args() num_workers = args.batch_size if args.batch_size <= 16 else 16 video_file = args.video_file video_name = video_file.split('/')[-1][:-4] video_range = args.video_range if args.video_range else None large_video = args.large_video out_csv_file = os.path.join(args.save_dir, f'{video_name}_ball.csv') out_video_file = os.path.join(args.save_dir, f'{video_name}.mp4') if not os.path.exists(args.save_dir): os.makedirs(args.save_dir) # Load model tracknet_ckpt = torch.load(args.tracknet_file) tracknet_seq_len = tracknet_ckpt['param_dict']['seq_len'] bg_mode = tracknet_ckpt['param_dict']['bg_mode'] tracknet = get_model('TrackNet', tracknet_seq_len, bg_mode).cuda() tracknet.load_state_dict(tracknet_ckpt['model']) if args.inpaintnet_file: inpaintnet_ckpt = torch.load(args.inpaintnet_file) inpaintnet_seq_len = inpaintnet_ckpt['param_dict']['seq_len'] inpaintnet = get_model('InpaintNet').cuda() inpaintnet.load_state_dict(inpaintnet_ckpt['model']) else: inpaintnet = None cap = cv2.VideoCapture(args.video_file) w, h = (int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)), int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))) w_scaler, h_scaler = w / WIDTH, h / HEIGHT img_scaler = (w_scaler, h_scaler) tracknet_pred_dict = {'Frame':[], 'X':[], 'Y':[], 'Visibility':[], 'Inpaint_Mask':[], 'Img_scaler': (w_scaler, h_scaler), 'Img_shape': (w, h)} # Test on TrackNet tracknet.eval() seq_len = tracknet_seq_len if args.eval_mode == 'nonoverlap': # Create dataset with non-overlap sampling if large_video: dataset = Video_IterableDataset(video_file, seq_len=seq_len, sliding_step=seq_len, bg_mode=bg_mode, max_sample_num=args.max_sample_num, video_range=video_range) data_loader = DataLoader(dataset, batch_size=args.batch_size, shuffle=False, drop_last=False) print(f'Video length: {dataset.video_len}') else: # Sample all frames from video frame_list = generate_frames(args.video_file) dataset = Shuttlecock_Trajectory_Dataset(seq_len=seq_len, sliding_step=seq_len, data_mode='heatmap', bg_mode=bg_mode, frame_arr=np.array(frame_list)[:, :, :, ::-1], padding=True) data_loader = DataLoader(dataset, batch_size=args.batch_size, shuffle=False, num_workers=num_workers, drop_last=False) for step, (i, x) in enumerate(tqdm(data_loader)): x = x.float().cuda() with torch.no_grad(): y_pred = tracknet(x).detach().cpu() # Predict tmp_pred = predict(i, y_pred=y_pred, img_scaler=img_scaler) for key in tmp_pred.keys(): tracknet_pred_dict[key].extend(tmp_pred[key]) else: # Create dataset with overlap sampling for temporal ensemble if large_video: dataset = Video_IterableDataset(video_file, seq_len=seq_len, sliding_step=1, bg_mode=bg_mode, max_sample_num=args.max_sample_num, video_range=video_range) data_loader = DataLoader(dataset, batch_size=args.batch_size, shuffle=False, drop_last=False) video_len = dataset.video_len print(f'Video length: {video_len}') else: # Sample all frames from video frame_list = generate_frames(args.video_file) dataset = Shuttlecock_Trajectory_Dataset(seq_len=seq_len, sliding_step=1, data_mode='heatmap', bg_mode=bg_mode, frame_arr=np.array(frame_list)[:, :, :, ::-1]) data_loader = DataLoader(dataset, batch_size=args.batch_size, shuffle=False, num_workers=num_workers, drop_last=False) video_len = len(frame_list) # Init prediction buffer params num_sample, sample_count = video_len-seq_len+1, 0 buffer_size = seq_len - 1 batch_i = torch.arange(seq_len) # [0, 1, 2, 3, 4, 5, 6, 7] frame_i = torch.arange(seq_len-1, -1, -1) # [7, 6, 5, 4, 3, 2, 1, 0] y_pred_buffer = torch.zeros((buffer_size, seq_len, HEIGHT, WIDTH), dtype=torch.float32) weight = get_ensemble_weight(seq_len, args.eval_mode) for step, (i, x) in enumerate(tqdm(data_loader)): x = x.float().cuda() b_size, seq_len = i.shape[0], i.shape[1] with torch.no_grad(): y_pred = tracknet(x).detach().cpu() y_pred_buffer = torch.cat((y_pred_buffer, y_pred), dim=0) ensemble_i = torch.empty((0, 1, 2), dtype=torch.float32) ensemble_y_pred = torch.empty((0, 1, HEIGHT, WIDTH), dtype=torch.float32) for b in range(b_size): if sample_count < buffer_size: # Imcomplete buffer y_pred = y_pred_buffer[batch_i+b, frame_i].sum(0) / (sample_count+1) else: # General case y_pred = (y_pred_buffer[batch_i+b, frame_i] * weight[:, None, None]).sum(0) ensemble_i = torch.cat((ensemble_i, i[b][0].reshape(1, 1, 2)), dim=0) ensemble_y_pred = torch.cat((ensemble_y_pred, y_pred.reshape(1, 1, HEIGHT, WIDTH)), dim=0) sample_count += 1 if sample_count == num_sample: # Last batch y_zero_pad = torch.zeros((buffer_size, seq_len, HEIGHT, WIDTH), dtype=torch.float32) y_pred_buffer = torch.cat((y_pred_buffer, y_zero_pad), dim=0) for f in range(1, seq_len): # Last input sequence y_pred = y_pred_buffer[batch_i+b+f, frame_i].sum(0) / (seq_len-f) ensemble_i = torch.cat((ensemble_i, i[-1][f].reshape(1, 1, 2)), dim=0) ensemble_y_pred = torch.cat((ensemble_y_pred, y_pred.reshape(1, 1, HEIGHT, WIDTH)), dim=0) # Predict tmp_pred = predict(ensemble_i, y_pred=ensemble_y_pred, img_scaler=img_scaler) for key in tmp_pred.keys(): tracknet_pred_dict[key].extend(tmp_pred[key]) # Update buffer, keep last predictions for ensemble in next iteration y_pred_buffer = y_pred_buffer[-buffer_size:] #assert video_len == len(tracknet_pred_dict['Frame']), 'Prediction length mismatch' # Test on TrackNetV3 (TrackNet + InpaintNet) if inpaintnet is not None: inpaintnet.eval() seq_len = inpaintnet_seq_len tracknet_pred_dict['Inpaint_Mask'] = generate_inpaint_mask(tracknet_pred_dict, th_h=h*0.05) inpaint_pred_dict = {'Frame':[], 'X':[], 'Y':[], 'Visibility':[]} if args.eval_mode == 'nonoverlap': # Create dataset with non-overlap sampling dataset = Shuttlecock_Trajectory_Dataset(seq_len=seq_len, sliding_step=seq_len, data_mode='coordinate', pred_dict=tracknet_pred_dict, padding=True) data_loader = DataLoader(dataset, batch_size=args.batch_size, shuffle=False, num_workers=num_workers, drop_last=False) for step, (i, coor_pred, inpaint_mask) in enumerate(tqdm(data_loader)): coor_pred, inpaint_mask = coor_pred.float(), inpaint_mask.float() with torch.no_grad(): coor_inpaint = inpaintnet(coor_pred.cuda(), inpaint_mask.cuda()).detach().cpu() coor_inpaint = coor_inpaint * inpaint_mask + coor_pred * (1-inpaint_mask) # replace predicted coordinates with inpainted coordinates # Thresholding th_mask = ((coor_inpaint[:, :, 0] < COOR_TH) & (coor_inpaint[:, :, 1] < COOR_TH)) coor_inpaint[th_mask] = 0. # Predict tmp_pred = predict(i, c_pred=coor_inpaint, img_scaler=img_scaler) for key in tmp_pred.keys(): inpaint_pred_dict[key].extend(tmp_pred[key]) else: # Create dataset with overlap sampling for temporal ensemble dataset = Shuttlecock_Trajectory_Dataset(seq_len=seq_len, sliding_step=1, data_mode='coordinate', pred_dict=tracknet_pred_dict) data_loader = DataLoader(dataset, batch_size=args.batch_size, shuffle=False, num_workers=num_workers, drop_last=False) weight = get_ensemble_weight(seq_len, args.eval_mode) # Init buffer params num_sample, sample_count = len(dataset), 0 buffer_size = seq_len - 1 batch_i = torch.arange(seq_len) # [0, 1, 2, 3, 4, 5, 6, 7] frame_i = torch.arange(seq_len-1, -1, -1) # [7, 6, 5, 4, 3, 2, 1, 0] coor_inpaint_buffer = torch.zeros((buffer_size, seq_len, 2), dtype=torch.float32) for step, (i, coor_pred, inpaint_mask) in enumerate(tqdm(data_loader)): coor_pred, inpaint_mask = coor_pred.float(), inpaint_mask.float() b_size = i.shape[0] with torch.no_grad(): coor_inpaint = inpaintnet(coor_pred.cuda(), inpaint_mask.cuda()).detach().cpu() coor_inpaint = coor_inpaint * inpaint_mask + coor_pred * (1-inpaint_mask) # Thresholding th_mask = ((coor_inpaint[:, :, 0] < COOR_TH) & (coor_inpaint[:, :, 1] < COOR_TH)) coor_inpaint[th_mask] = 0. coor_inpaint_buffer = torch.cat((coor_inpaint_buffer, coor_inpaint), dim=0) ensemble_i = torch.empty((0, 1, 2), dtype=torch.float32) ensemble_coor_inpaint = torch.empty((0, 1, 2), dtype=torch.float32) for b in range(b_size): if sample_count < buffer_size: # Imcomplete buffer coor_inpaint = coor_inpaint_buffer[batch_i+b, frame_i].sum(0) coor_inpaint /= (sample_count+1) else: # General case coor_inpaint = (coor_inpaint_buffer[batch_i+b, frame_i] * weight[:, None]).sum(0) ensemble_i = torch.cat((ensemble_i, i[b][0].view(1, 1, 2)), dim=0) ensemble_coor_inpaint = torch.cat((ensemble_coor_inpaint, coor_inpaint.view(1, 1, 2)), dim=0) sample_count += 1 if sample_count == num_sample: # Last input sequence coor_zero_pad = torch.zeros((buffer_size, seq_len, 2), dtype=torch.float32) coor_inpaint_buffer = torch.cat((coor_inpaint_buffer, coor_zero_pad), dim=0) for f in range(1, seq_len): coor_inpaint = coor_inpaint_buffer[batch_i+b+f, frame_i].sum(0) coor_inpaint /= (seq_len-f) ensemble_i = torch.cat((ensemble_i, i[-1][f].view(1, 1, 2)), dim=0) ensemble_coor_inpaint = torch.cat((ensemble_coor_inpaint, coor_inpaint.view(1, 1, 2)), dim=0) # Thresholding th_mask = ((ensemble_coor_inpaint[:, :, 0] < COOR_TH) & (ensemble_coor_inpaint[:, :, 1] < COOR_TH)) ensemble_coor_inpaint[th_mask] = 0. # Predict tmp_pred = predict(ensemble_i, c_pred=ensemble_coor_inpaint, img_scaler=img_scaler) for key in tmp_pred.keys(): inpaint_pred_dict[key].extend(tmp_pred[key]) # Update buffer, keep last predictions for ensemble in next iteration coor_inpaint_buffer = coor_inpaint_buffer[-buffer_size:] # Write csv file pred_dict = inpaint_pred_dict if inpaintnet is not None else tracknet_pred_dict write_pred_csv(pred_dict, save_file=out_csv_file) # Write video with predicted coordinates if args.output_video: write_pred_video(video_file, pred_dict, save_file=out_video_file, traj_len=args.traj_len) print('Done.')

csv进行可视化；

import cv2 import pandas as pd import numpy as np import os def draw_trajectory_on_video(video_path, csv_path, output_path, traj_len=8): """在视频上绘制轨迹""" # 读取CSV数据 df = pd.read_csv(csv_path) # 如果CSV没有表头，使用列名 if 'Frame' not in df.columns: df.columns = ['Frame', 'Visibility', 'X', 'Y'] # 打开视频 cap = cv2.VideoCapture(video_path) fps = int(cap.get(cv2.CAP_PROP_FPS)) width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) # 视频写入器 fourcc = cv2.VideoWriter_fourcc(*'mp4v') out = cv2.VideoWriter(output_path, fourcc, fps, (width, height)) frame_idx = 0 positions = [] # 存储最近的轨迹点 while True: ret, frame = cap.read() if not ret: break # 获取当前帧的预测数据 row = df[df['Frame'] == frame_idx] if not row.empty: x = int(row['X'].values[0]) y = int(row['Y'].values[0]) vis = int(row['Visibility'].values[0]) if vis == 1 and x > 0 and y > 0: # 添加当前点 positions.append((x, y)) # 只保留最近traj_len个点 if len(positions) > traj_len: positions.pop(0) # 绘制轨迹（从旧到新） for i in range(1, len(positions)): # 颜色从蓝色渐变到红色 color = (0, int(255 * i / len(positions)), int(255 * (1 - i / len(positions)))) cv2.line(frame, positions[ i -1], positions[i], color, 1) # 绘制当前点（大圆圈） cv2.circle(frame, (x, y), 6, (0, 0, 255), 1) cv2.circle(frame, (x, y), 8, (255, 255, 255), 1) # 显示坐标信息 cv2.putText(frame, f'({x}, {y})', ( x +15, y- 15), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 255), 2) # 显示帧号 cv2.putText(frame, f'Frame: {frame_idx}', (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 255, 255), 2) out.write(frame) cv2.imshow('frame', frame) cv2.waitKey(0) frame_idx += 1 # 进度显示 if frame_idx % 100 == 0: print(f'Processing frame {frame_idx}/{total_frames}') cap.release() out.release() cv2.destroyAllWindows() print(f'视频已保存到: {output_path}') if __name__ == '__main__': # 请修改这些路径 video_path = "C:/Users/ChanJing-01/Videos/yumao/yumao1.mp4" csv_path = r"C:/Users/ChanJing-01/Videos/yumao/yumao1_ball.csv" # 您的CSV文件路径 output_path = "pred_result/yumao1_with_trajectory.mp4" draw_trajectory_on_video(video_path, csv_path, output_path, traj_len=8)