目录极小目标检测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_predNone, c_predNone, 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, typestr, defaultrC:\Users\ChanJing-01\Videos\yumao\yumao1.mp4, helpfile path of the video) parser.add_argument(--tracknet_file, typestr, defaultrD:\project_2026\yumaoqiu\TrackNetV3-master\ckpts\TrackNet_best.pt, helpfile path of the TrackNet model checkpoint) parser.add_argument(--inpaintnet_file, typestr, defaultrD:\project_2026\yumaoqiu\TrackNetV3-master\ckpts\InpaintNet_best.pt, helpfile path of the InpaintNet model checkpoint) # parser.add_argument(--video_file, typestr, helpfile path of the video) # parser.add_argument(--tracknet_file, typestr, helpfile path of the TrackNet model checkpoint) # parser.add_argument(--inpaintnet_file, typestr, default, helpfile path of the InpaintNet model checkpoint) parser.add_argument(--batch_size, typeint, default16, helpbatch size for inference) parser.add_argument(--eval_mode, typestr, defaultweight, choices[nonoverlap, average, weight], helpevaluation mode) parser.add_argument(--max_sample_num, typeint, default1800, helpmaximum number of frames to sample for generating median image) parser.add_argument(--video_range, typelambda splits: [int(s) for s in splits.split(,)], defaultNone, helprange of start second and end second of the video for generating median image) parser.add_argument(--save_dir, typestr, defaultpred_result, helpdirectory to save the prediction result) parser.add_argument(--large_video, actionstore_true, defaultTrue, helpwhether to process large video) parser.add_argument(--output_video, actionstore_true, defaultTrue, helpwhether to output video with predicted trajectory) parser.add_argument(--traj_len, typeint, default8, helplength 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_lenseq_len, sliding_stepseq_len, bg_modebg_mode, max_sample_numargs.max_sample_num, video_rangevideo_range) data_loader DataLoader(dataset, batch_sizeargs.batch_size, shuffleFalse, drop_lastFalse) print(fVideo length: {dataset.video_len}) else: # Sample all frames from video frame_list generate_frames(args.video_file) dataset Shuttlecock_Trajectory_Dataset(seq_lenseq_len, sliding_stepseq_len, data_modeheatmap, bg_modebg_mode, frame_arrnp.array(frame_list)[:, :, :, ::-1], paddingTrue) data_loader DataLoader(dataset, batch_sizeargs.batch_size, shuffleFalse, num_workersnum_workers, drop_lastFalse) 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_predy_pred, img_scalerimg_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_lenseq_len, sliding_step1, bg_modebg_mode, max_sample_numargs.max_sample_num, video_rangevideo_range) data_loader DataLoader(dataset, batch_sizeargs.batch_size, shuffleFalse, drop_lastFalse) video_len dataset.video_len print(fVideo length: {video_len}) else: # Sample all frames from video frame_list generate_frames(args.video_file) dataset Shuttlecock_Trajectory_Dataset(seq_lenseq_len, sliding_step1, data_modeheatmap, bg_modebg_mode, frame_arrnp.array(frame_list)[:, :, :, ::-1]) data_loader DataLoader(dataset, batch_sizeargs.batch_size, shuffleFalse, num_workersnum_workers, drop_lastFalse) video_len len(frame_list) # Init prediction buffer params num_sample, sample_count video_len-seq_len1, 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), dtypetorch.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), dim0) ensemble_i torch.empty((0, 1, 2), dtypetorch.float32) ensemble_y_pred torch.empty((0, 1, HEIGHT, WIDTH), dtypetorch.float32) for b in range(b_size): if sample_count buffer_size: # Imcomplete buffer y_pred y_pred_buffer[batch_ib, frame_i].sum(0) / (sample_count1) else: # General case y_pred (y_pred_buffer[batch_ib, frame_i] * weight[:, None, None]).sum(0) ensemble_i torch.cat((ensemble_i, i[b][0].reshape(1, 1, 2)), dim0) ensemble_y_pred torch.cat((ensemble_y_pred, y_pred.reshape(1, 1, HEIGHT, WIDTH)), dim0) sample_count 1 if sample_count num_sample: # Last batch y_zero_pad torch.zeros((buffer_size, seq_len, HEIGHT, WIDTH), dtypetorch.float32) y_pred_buffer torch.cat((y_pred_buffer, y_zero_pad), dim0) for f in range(1, seq_len): # Last input sequence y_pred y_pred_buffer[batch_ibf, frame_i].sum(0) / (seq_len-f) ensemble_i torch.cat((ensemble_i, i[-1][f].reshape(1, 1, 2)), dim0) ensemble_y_pred torch.cat((ensemble_y_pred, y_pred.reshape(1, 1, HEIGHT, WIDTH)), dim0) # Predict tmp_pred predict(ensemble_i, y_predensemble_y_pred, img_scalerimg_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_hh*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_lenseq_len, sliding_stepseq_len, data_modecoordinate, pred_dicttracknet_pred_dict, paddingTrue) data_loader DataLoader(dataset, batch_sizeargs.batch_size, shuffleFalse, num_workersnum_workers, drop_lastFalse) 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_predcoor_inpaint, img_scalerimg_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_lenseq_len, sliding_step1, data_modecoordinate, pred_dicttracknet_pred_dict) data_loader DataLoader(dataset, batch_sizeargs.batch_size, shuffleFalse, num_workersnum_workers, drop_lastFalse) 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), dtypetorch.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), dim0) ensemble_i torch.empty((0, 1, 2), dtypetorch.float32) ensemble_coor_inpaint torch.empty((0, 1, 2), dtypetorch.float32) for b in range(b_size): if sample_count buffer_size: # Imcomplete buffer coor_inpaint coor_inpaint_buffer[batch_ib, frame_i].sum(0) coor_inpaint / (sample_count1) else: # General case coor_inpaint (coor_inpaint_buffer[batch_ib, frame_i] * weight[:, None]).sum(0) ensemble_i torch.cat((ensemble_i, i[b][0].view(1, 1, 2)), dim0) ensemble_coor_inpaint torch.cat((ensemble_coor_inpaint, coor_inpaint.view(1, 1, 2)), dim0) sample_count 1 if sample_count num_sample: # Last input sequence coor_zero_pad torch.zeros((buffer_size, seq_len, 2), dtypetorch.float32) coor_inpaint_buffer torch.cat((coor_inpaint_buffer, coor_zero_pad), dim0) for f in range(1, seq_len): coor_inpaint coor_inpaint_buffer[batch_ibf, frame_i].sum(0) coor_inpaint / (seq_len-f) ensemble_i torch.cat((ensemble_i, i[-1][f].view(1, 1, 2)), dim0) ensemble_coor_inpaint torch.cat((ensemble_coor_inpaint, coor_inpaint.view(1, 1, 2)), dim0) # 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_predensemble_coor_inpaint, img_scalerimg_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_fileout_csv_file) # Write video with predicted coordinates if args.output_video: write_pred_video(video_file, pred_dict, save_fileout_video_file, traj_lenargs.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_len8): 在视频上绘制轨迹 # 读取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, fFrame: {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(fProcessing 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 rC:/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_len8)