diff --git a/.DS_Store b/.DS_Store new file mode 100644 index 0000000..9bfdabc Binary files /dev/null and b/.DS_Store differ diff --git a/.gitignore b/.gitignore index 076a0d4..898ef24 100644 --- a/.gitignore +++ b/.gitignore @@ -6,6 +6,7 @@ figure/ __pycache__/ utils/__pycache__/ prediction/ +videos/ test_prediction/ dataset_test.py fps_exp.py @@ -15,4 +16,4 @@ temp.py model_compare.py predict_old.py *.sh -*.mp4 \ No newline at end of file +*.mp4 diff --git a/README.md b/README.md index 9d168e2..04736c1 100644 --- a/README.md +++ b/README.md @@ -77,6 +77,13 @@ scenarios to strengthen the network’s robustness. Given that a shuttlecock can 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 ``` +### Experiment logging +* `predict.py` appends a row to `experiments/experiment_log.csv` after each successful inference run. + ``` + python predict.py --video_file videos/test-okimoto.mp4 --tracknet_file ckpts/TrackNet_best.pt --inpaintnet_file ckpts/InpaintNet_best.pt --save_dir prediction + ``` +* Use `--no_experiment_log` to disable automatic logging for a run. + ## Training ### 1. Prepare Dataset * Download [Shuttlecock Trajectory Dataset](https://hackmd.io/Nf8Rh1NrSrqNUzmO0sQKZw) diff --git a/dataset.py b/dataset.py index a33a18d..9a69fc9 100644 --- a/dataset.py +++ b/dataset.py @@ -765,8 +765,9 @@ def __gen_median__(self, max_sample_num, video_range): else: sample_step = 1 + frame_indices = range(start_frame, end_frame, sample_step) frame_list = [] - for i in range(start_frame, end_frame, sample_step): + for i in tqdm(frame_indices, desc='median_s'): self.cap.set(cv2.CAP_PROP_POS_FRAMES, i) success, frame = self.cap.read() if not success: @@ -811,4 +812,4 @@ def __process__(self, imgs): frames /= 255. return frames - \ No newline at end of file + diff --git a/experiments/experiment_log.csv b/experiments/experiment_log.csv new file mode 100644 index 0000000..2f1246a --- /dev/null +++ b/experiments/experiment_log.csv @@ -0,0 +1,9 @@ +date,id,video_info,angle,background,preprocessing,resolution,fps,duration_s,bitrate,runtime_s,model_load_s,median_s,inference_s,write_csv_s,video_merge_s,result_notes,input_video,memo,accuracy,command +2026-06-27,1,BWFの試合映像 沖本,俯瞰,グリーンマット,なし,854x480,24,55,478421,,,,,,,,test-okimoto.mp4,,, +2026-06-27,3,全国選抜2024 山田,俯瞰,普通のコート,なし,1920x1080,60,40,14062749,1164,,,,,,,yamada-senbatu-crop.mp4,,, +2026-06-27,4,全国選抜2024山田,俯瞰,普通のコート,あり,1280x720,30,40,968864,,,,,,,,formated-yamada-senbatu-crop-720p30.mp4,,, +2026-06-28,5,国体2025 長束,右斜め上,グリーンマット,あり,1280x720,30,30,2825850,,,,,,,,,, +2026-06-28,6,総合2016 エンワタ,ギリ俯瞰,グリーンマット,あり,1280x720,30,93,1149310,,,,,,,,,, +2026-06-28,7,神々の遊び 3v3,横、手ブレ,選手の雑踏,あり,1280x720,30,30,2825850,,,,,,,,,, +2026-06-28,8,宮大体育館 辰樹史弥,左下,体育館の風景,あり,1280x720,30,358,1328318,,,,,,,,,, +1782803569.0459,9,,,,,1920x1080,25,46.16,2057342,844.959,0.287,272.909,559.084,0.051,12.467,,okimoto-us.mp4,,,predict.py --video_file videos/okimoto-us.mp4 --tracknet_file ckpts/TrackNet_best.pt --inpaintnet_file ckpts/InpaintNet_best.pt --save_dir prediction --batch_size 4 --output_video --large_video diff --git a/predict.py b/predict.py index 82aa285..0098b2f 100644 --- a/predict.py +++ b/predict.py @@ -1,6 +1,9 @@ import os import argparse +import sys +import time import numpy as np +import cv2 from tqdm import tqdm import torch @@ -8,6 +11,7 @@ from test import predict_location, get_ensemble_weight, generate_inpaint_mask from dataset import Shuttlecock_Trajectory_Dataset, Video_IterableDataset +from utils.experiment_log import append_prediction_log from utils.general import * @@ -81,7 +85,16 @@ def predict(indices, y_pred=None, c_pred=None, img_scaler=(1, 1)): parser.add_argument('--large_video', action='store_true', default=False, help='whether to process large video') parser.add_argument('--output_video', action='store_true', default=False, 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') + parser.add_argument('--no_experiment_log', action='store_true', default=False, help='disable automatic experiment logging') args = parser.parse_args() + run_start_time = time.perf_counter() + stage_times = { + 'model_load_s': 0.0, + 'median_s': 0.0, + 'inference_s': 0.0, + 'write_csv_s': 0.0, + 'video_merge_s': 0.0, + } num_workers = args.batch_size if args.batch_size <= 16 else 16 video_file = args.video_file @@ -95,22 +108,27 @@ def predict(indices, y_pred=None, c_pred=None, img_scaler=(1, 1)): os.makedirs(args.save_dir) # Load model - tracknet_ckpt = torch.load(args.tracknet_file) + stage_start_time = time.perf_counter() + device = torch.device("mps" if torch.backends.mps.is_available() else "cpu") + + tracknet_ckpt = torch.load(args.tracknet_file, map_location=device) 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 = get_model('TrackNet', tracknet_seq_len, bg_mode).to(device) tracknet.load_state_dict(tracknet_ckpt['model']) if args.inpaintnet_file: - inpaintnet_ckpt = torch.load(args.inpaintnet_file) + inpaintnet_ckpt = torch.load(args.inpaintnet_file, map_location=device) inpaintnet_seq_len = inpaintnet_ckpt['param_dict']['seq_len'] - inpaintnet = get_model('InpaintNet').cuda() + inpaintnet = get_model('InpaintNet').to(device) inpaintnet.load_state_dict(inpaintnet_ckpt['model']) else: inpaintnet = None + stage_times['model_load_s'] += time.perf_counter() - stage_start_time cap = cv2.VideoCapture(args.video_file) w, h = (int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)), int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))) + cap.release() w_scaler, h_scaler = w / WIDTH, h / HEIGHT img_scaler = (w_scaler, h_scaler) @@ -122,8 +140,9 @@ def predict(indices, y_pred=None, c_pred=None, img_scaler=(1, 1)): seq_len = tracknet_seq_len if args.eval_mode == 'nonoverlap': # Create dataset with non-overlap sampling + stage_start_time = time.perf_counter() if large_video: - dataset = Video_IterableDataset(video_file, seq_len=seq_len, sliding_step=seq_len, bg_mode=bg_mode, + 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}') @@ -133,9 +152,11 @@ def predict(indices, y_pred=None, c_pred=None, img_scaler=(1, 1)): 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) + stage_times['median_s'] += time.perf_counter() - stage_start_time - for step, (i, x) in enumerate(tqdm(data_loader)): - x = x.float().cuda() + stage_start_time = time.perf_counter() + for step, (i, x) in enumerate(tqdm(data_loader, desc='tracknet_inference_s')): + x = x.float().to(device) with torch.no_grad(): y_pred = tracknet(x).detach().cpu() @@ -143,8 +164,10 @@ def predict(indices, y_pred=None, c_pred=None, img_scaler=(1, 1)): 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]) + stage_times['inference_s'] += time.perf_counter() - stage_start_time else: # Create dataset with overlap sampling for temporal ensemble + stage_start_time = time.perf_counter() 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) @@ -159,6 +182,7 @@ def predict(indices, y_pred=None, c_pred=None, img_scaler=(1, 1)): 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) + stage_times['median_s'] += time.perf_counter() - stage_start_time # Init prediction buffer params num_sample, sample_count = video_len-seq_len+1, 0 @@ -167,8 +191,9 @@ def predict(indices, y_pred=None, c_pred=None, img_scaler=(1, 1)): 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() + stage_start_time = time.perf_counter() + for step, (i, x) in enumerate(tqdm(data_loader, desc='tracknet_inference_s')): + x = x.float().to(device) b_size, seq_len = i.shape[0], i.shape[1] with torch.no_grad(): y_pred = tracknet(x).detach().cpu() @@ -207,6 +232,7 @@ def predict(indices, y_pred=None, c_pred=None, img_scaler=(1, 1)): # Update buffer, keep last predictions for ensemble in next iteration y_pred_buffer = y_pred_buffer[-buffer_size:] + stage_times['inference_s'] += time.perf_counter() - stage_start_time #assert video_len == len(tracknet_pred_dict['Frame']), 'Prediction length mismatch' # Test on TrackNetV3 (TrackNet + InpaintNet) @@ -221,10 +247,11 @@ def predict(indices, y_pred=None, c_pred=None, img_scaler=(1, 1)): 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)): + stage_start_time = time.perf_counter() + for step, (i, coor_pred, inpaint_mask) in enumerate(tqdm(data_loader, desc='inpaintnet_inference_s')): 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 = inpaintnet(coor_pred.to(device), inpaint_mask.to(device)).detach().cpu() coor_inpaint = coor_inpaint * inpaint_mask + coor_pred * (1-inpaint_mask) # replace predicted coordinates with inpainted coordinates # Thresholding @@ -235,6 +262,7 @@ def predict(indices, y_pred=None, c_pred=None, img_scaler=(1, 1)): 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]) + stage_times['inference_s'] += time.perf_counter() - stage_start_time else: # Create dataset with overlap sampling for temporal ensemble @@ -249,11 +277,12 @@ def predict(indices, y_pred=None, c_pred=None, img_scaler=(1, 1)): 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)): + stage_start_time = time.perf_counter() + for step, (i, coor_pred, inpaint_mask) in enumerate(tqdm(data_loader, desc='inpaintnet_inference_s')): 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 = inpaintnet(coor_pred.to(device), inpaint_mask.to(device)).detach().cpu() coor_inpaint = coor_inpaint * inpaint_mask + coor_pred * (1-inpaint_mask) # Thresholding @@ -299,14 +328,24 @@ def predict(indices, y_pred=None, c_pred=None, img_scaler=(1, 1)): # Update buffer, keep last predictions for ensemble in next iteration coor_inpaint_buffer = coor_inpaint_buffer[-buffer_size:] + stage_times['inference_s'] += time.perf_counter() - stage_start_time # Write csv file pred_dict = inpaint_pred_dict if inpaintnet is not None else tracknet_pred_dict + stage_start_time = time.perf_counter() write_pred_csv(pred_dict, save_file=out_csv_file) + stage_times['write_csv_s'] += time.perf_counter() - stage_start_time # Write video with predicted coordinates if args.output_video: + stage_start_time = time.perf_counter() write_pred_video(video_file, pred_dict, save_file=out_video_file, traj_len=args.traj_len) + stage_times['video_merge_s'] += time.perf_counter() - stage_start_time + + runtime_s = time.perf_counter() - run_start_time + + if not args.no_experiment_log: + append_prediction_log(args, runtime_s, stage_times, sys.argv) - print('Done.') \ No newline at end of file + print('Done.') diff --git a/requirements.txt b/requirements.txt index 9d5d87c..47b4380 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,8 +1,10 @@ dash==2.5.1 numpy==1.22.4 -opencv_python==4.4.0.46 +opencv_python==4.13.0.92 pandas==2.0.0 Pillow==10.0.0 plotly==5.8.2 -torch==1.10.0 -parse \ No newline at end of file +torch==2.4.1 +parse +tqdm==4.68.3 +pycocotools==2.0.7 diff --git a/utils/experiment_log.py b/utils/experiment_log.py new file mode 100644 index 0000000..7a3e11c --- /dev/null +++ b/utils/experiment_log.py @@ -0,0 +1,122 @@ +import csv +import os +import shlex +import time +from pathlib import Path + +import cv2 + + +EXPERIMENT_LOG_FILE = Path("experiments/experiment_log.csv") +EXPERIMENT_LOG_FIELDS = [ + "date", + "id", + "video_info", + "angle", + "background", + "preprocessing", + "resolution", + "fps", + "duration_s", + "bitrate", + "runtime_s", + "model_load_s", + "median_s", + "inference_s", + "write_csv_s", + "video_merge_s", + "result_notes", + "input_video", + "memo", + "accuracy", + "command", +] + + +def format_seconds(seconds): + return f"{seconds:.3f}" + + +def read_experiment_rows(log_file): + if not log_file.exists(): + return [] + + with log_file.open("r", encoding="utf-8", newline="") as f: + return list(csv.DictReader(f)) + + +def next_experiment_id(rows): + ids = [] + for row in rows: + try: + ids.append(int(row.get("id", ""))) + except ValueError: + continue + return max(ids, default=0) + 1 + + +def get_video_metadata(video_file): + cap = cv2.VideoCapture(video_file) + if not cap.isOpened(): + return { + "resolution": "", + "fps": "", + "duration_s": "", + "bitrate": "", + } + + width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) + height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) + fps = cap.get(cv2.CAP_PROP_FPS) + frame_count = cap.get(cv2.CAP_PROP_FRAME_COUNT) + cap.release() + + duration_s = frame_count / fps if fps and fps > 0 else 0 + file_size_bits = os.path.getsize(video_file) * 8 if os.path.exists(video_file) else 0 + bitrate = int(file_size_bits / duration_s) if duration_s > 0 else "" + + return { + "resolution": f"{width}x{height}" if width and height else "", + "fps": f"{fps:.3f}".rstrip("0").rstrip(".") if fps else "", + "duration_s": f"{duration_s:.3f}".rstrip("0").rstrip(".") if duration_s else "", + "bitrate": bitrate, + } + + +def append_prediction_log(args, runtime_s, stage_times, argv): + log_file = EXPERIMENT_LOG_FILE + metadata = get_video_metadata(args.video_file) + rows = read_experiment_rows(log_file) + row = { + "date": time.time(), + "id": next_experiment_id(rows), + "video_info": "", + "angle": "", + "background": "", + "preprocessing": "", + "resolution": metadata["resolution"], + "fps": metadata["fps"], + "duration_s": metadata["duration_s"], + "bitrate": metadata["bitrate"], + "runtime_s": format_seconds(runtime_s), + "model_load_s": format_seconds(stage_times.get('model_load_s', 0)), + "median_s": format_seconds(stage_times.get('median_s', 0)), + "inference_s": format_seconds(stage_times.get('inference_s', 0)), + "write_csv_s": format_seconds(stage_times.get('write_csv_s', 0)), + "video_merge_s": format_seconds(stage_times.get('video_merge_s', 0)), + "result_notes": "", + "input_video": os.path.basename(args.video_file), + "memo": "", + "accuracy": "", + "command": " ".join(shlex.quote(arg) for arg in argv), + } + + log_file.parent.mkdir(parents=True, exist_ok=True) + exists = log_file.exists() + with log_file.open("a", encoding="utf-8", newline="") as f: + writer = csv.DictWriter(f, fieldnames=EXPERIMENT_LOG_FIELDS) + if not exists: + writer.writeheader() + writer.writerow(row) + + print(f"Experiment log: appended #{row['id']} to {log_file}") diff --git a/utils/general.py b/utils/general.py index 33213cb..f4a513a 100644 --- a/utils/general.py +++ b/utils/general.py @@ -9,6 +9,7 @@ from collections import deque from PIL import Image, ImageDraw +from tqdm import tqdm from model import TrackNet, InpaintNet # Global variables @@ -213,14 +214,18 @@ def generate_frames(video_file): # Get camera parameters cap = cv2.VideoCapture(video_file) + frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) frame_list = [] success = True # Sample frames until video end - while success: - success, frame = cap.read() - if success: - frame_list.append(frame) + with tqdm(total=frame_count if frame_count > 0 else None, desc='generate_frames_s') as pbar: + while success: + success, frame = cap.read() + if success: + frame_list.append(frame) + pbar.update(1) + cap.release() return frame_list @@ -290,31 +295,34 @@ def write_pred_video(video_file, pred_dict, save_file, traj_len=8, label_df=None # Draw label and prediction trajectory #for i, frame in enumerate(frame_list): i = 0 - while True: - success, frame = cap.read() - if not success: - break - - # Check capacity of queue - if len(pred_queue) >= traj_len: - pred_queue.pop() - if label_df is not None and len(gt_queue) >= traj_len: - gt_queue.pop() - - # Push ball coordinates for each frame - if label_df is not None: - gt_queue.appendleft([x[i], y[i]]) if vis[i] and i < len(label_df) else gt_queue.appendleft(None) - pred_queue.appendleft([x_pred[i], y_pred[i]]) if vis_pred[i] else pred_queue.appendleft(None) - - # Draw ground truth trajectory if exists - if label_df is not None: - frame = draw_traj(frame, gt_queue, color='red') - - # Draw prediction trajectory - frame = draw_traj(frame, pred_queue, color='yellow') + frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) + with tqdm(total=frame_count if frame_count > 0 else None, desc='video_merge_s') as pbar: + while True: + success, frame = cap.read() + if not success: + break + + # Check capacity of queue + if len(pred_queue) >= traj_len: + pred_queue.pop() + if label_df is not None and len(gt_queue) >= traj_len: + gt_queue.pop() + + # Push ball coordinates for each frame + if label_df is not None: + gt_queue.appendleft([x[i], y[i]]) if vis[i] and i < len(label_df) else gt_queue.appendleft(None) + pred_queue.appendleft([x_pred[i], y_pred[i]]) if vis_pred[i] else pred_queue.appendleft(None) + + # Draw ground truth trajectory if exists + if label_df is not None: + frame = draw_traj(frame, gt_queue, color='red') + + # Draw prediction trajectory + frame = draw_traj(frame, pred_queue, color='yellow') - out.write(frame) - i+=1 + out.write(frame) + i+=1 + pbar.update(1) out.release() cap.release() @@ -534,4 +542,4 @@ def re_generate_median_files(data_dir): print(f'Processing {video_file}...') get_rally_median(video_file) get_match_median(match_dir) - print(f'Finish processing {match_name}.') \ No newline at end of file + print(f'Finish processing {match_name}.')