Skip to content

Commit b7342e4

Browse files
fnunnarifnunnari
committed
Fixed frame extraction.
Implemented video rebuilding.
1 parent 810b63d commit b7342e4

File tree

2 files changed

+121
-25
lines changed

2 files changed

+121
-25
lines changed

‎PostProcessing/PostProcessVideos.py‎

Lines changed: 31 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
import argparse
22
from pathlib import Path
33
from typing import List, Tuple
4+
import tempfile
45

56
import pandas as pd
67
import re
@@ -9,9 +10,10 @@
910
from dataframes import repair_dropped_frames, compute_time_step
1011

1112
from video import extract_frames
13+
from video import rebuild_video
1214

1315

14-
THRESHOLD_NS = 10 * 1000 * 1000
16+
THRESHOLD_NS = 10 * 1000 * 1000 # millis * micros * nanos
1517

1618

1719
def scan_session_dir(input_dir: Path) -> Tuple[List[str], List[pd.DataFrame], List[str]]:
@@ -94,28 +96,40 @@ def main(input_dir: Path, output_dir: Path):
9496
# Find time ranges
9597
min_common, max_common = compute_time_range(repaired_df_list)
9698
# Trim the data frames to the time range
97-
trimmed_dataframes = trim_into_interval(repaired_df_list, min_common, max_common, THRESHOLD_NS)
99+
#trimmed_dataframes = trim_into_interval(repaired_df_list, min_common, max_common, THRESHOLD_NS)
100+
# TODO - temp to continue
101+
trimmed_dataframes = repaired_df_list
98102

99103
assert len(clientIDs) == len(trimmed_dataframes), f"Expected {len(clientIDs)} trimmed dataframes. Found f{len(trimmed_dataframes)}"
100104

101-
# Check that all the resultiong dataframes have the same number of rows
102-
client0ID = clientIDs[0]
103-
client0size = len(trimmed_dataframes[0])
104-
print(f"For client {client0ID}: {client0size} frames")
105-
for cID, df in zip(clientIDs[1:], trimmed_dataframes[1:]):
106-
dfsize = len(df)
107-
if client0size != dfsize:
108-
raise Exception(f"For client {cID}: expecting {client0size}, found {dfsize}")
105+
# Check that all the resulting dataframes have the same number of rows
106+
if False:
107+
client0ID = clientIDs[0]
108+
client0size = len(trimmed_dataframes[0])
109+
print(f"For client {client0ID}: {client0size} frames")
110+
for cID, df in zip(clientIDs[1:], trimmed_dataframes[1:]):
111+
dfsize = len(df)
112+
if client0size != dfsize:
113+
raise Exception(f"For client {cID}: expecting {client0size}, found {dfsize}")
109114

110-
print("Good. All trimmed dataframes have the same number of entries.")
115+
print("Good. All trimmed dataframes have the same number of entries.")
111116

112117
#
113-
# Extract the frames from the original videos
114-
# and rename the file names to the timestamps (DONE)
115-
#extract_frames(input_dir, output_dir)
116-
117-
#
118-
# Reconstruct videos (TODO)
118+
# Unpack the original videos, and repack them according to repaired and trimmed dataframes.
119+
for i, cID in enumerate(clientIDs):
120+
orig_df = df_list[i]
121+
trimmed_df = trimmed_dataframes[i]
122+
video_file = mp4_list[i]
123+
# Create a temporary directory for frames unpacking
124+
with tempfile.TemporaryDirectory(prefix="RecSyncNG", suffix=cID) as tmp_dir:
125+
# Extract the frames from the original videos
126+
# and rename the file names to the timestamps
127+
print(f"Extracting {len(orig_df)} frames from '{video_file}'...")
128+
extract_frames(video_file=video_file, timestamps_df=orig_df, output_dir=tmp_dir)
129+
130+
# Reconstruct videos (TODO)
131+
out_filepath = output_dir / (cID + ".mp4")
132+
rebuild_video(dir=Path(tmp_dir), frames=trimmed_df, outfile=out_filepath)
119133

120134

121135
#

‎PostProcessing/video.py‎

Lines changed: 90 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,9 @@
11
import cv2
2-
import os
2+
import os
3+
from pathlib import Path
34

45
import pandas as pd
6+
import numpy as np
57
import ffmpeg
68

79
from typing import Tuple
@@ -33,22 +35,102 @@ def video_info(video_path: str) -> Tuple[int, int, int]:
3335
return video_w, video_h, n_frames
3436

3537

36-
def extract_frames(video_file: str, timestamps: pd.DataFrame, output_dir: str):
38+
def extract_frames(video_file: str, timestamps_df: pd.DataFrame, output_dir: str):
39+
40+
#_, _, num_frames = video_info(video_file)
41+
#print("NUM:", num_frames, len(timestamps_df))
3742

3843
# Open the video file
3944
cap = cv2.VideoCapture(video_file)
40-
assert cap.isOpened() == True
41-
video_name, _ = os.path.splitext(video_file)
45+
if not cap.isOpened():
46+
raise Exception(f"Couldn't open video file '{video_file}'")
47+
48+
first_col_name = timestamps_df.columns[0]
49+
timestamps: pd.Series = timestamps_df[first_col_name]
4250

4351
# Loop over each timestamp in the CSV file
44-
for timestamp in timestamps:
52+
for fnum, timestamp in enumerate(timestamps):
4553
# Extract the frame using OpenCV
46-
cap.set(cv2.CAP_PROP_POS_MSEC, timestamp * 1000)
54+
#cap.set(cv2.CAP_PROP_POS_MSEC, timestamp * 1000)
4755
ret, frame = cap.read()
4856
if ret:
49-
frame_name = f"{video_name}_{timestamp:.3f}.jpg"
50-
frame_path = os.path.join(output_dir, frame_name)
57+
frame_path = os.path.join(output_dir, str(timestamp) + ".jpg")
5158
cv2.imwrite(frame_path, frame)
59+
else:
60+
print(f"At frame {fnum}, no more frames to extract from video '{video_file}'. Expected {len(timestamps)} frames.")
61+
#raise Exception(f"At frame {fnum}, no more frames to extract from video '{video_file}'. Expected {len(timestamps)} frames.")
5262

5363
# Release the video file
5464
cap.release()
65+
66+
67+
def rebuild_video(dir: Path, frames: pd.DataFrame, outfile: Path) -> None:
68+
69+
# We don't know the target video size, yet.
70+
frame_width = None
71+
frame_height = None
72+
73+
# It will be instantiated later, after we know the size of the first image
74+
ffmpeg_video_out_process = None
75+
76+
for idx, row in frames.iterrows():
77+
ts = row["timestamp"]
78+
gen = row["generated"]
79+
80+
if gen == "Original":
81+
82+
frame_path = dir / (str(ts) + ".jpg")
83+
84+
if not frame_path.exists():
85+
print(f"Skipping frame {str(frame_path)}")
86+
continue # BEWARE! Continues the cycle
87+
88+
img_bgr = cv2.imread(str(frame_path))
89+
img = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB)
90+
91+
# If the frame size was not determined yet, take it from the next picture and initialize the ffmpeg encoder
92+
if frame_width is None:
93+
assert frame_width is None and frame_height is None and ffmpeg_video_out_process is None
94+
95+
frame_height, frame_width, _ = img.shape
96+
font_width = int(frame_height * 0.04)
97+
ffmpeg_video_out_process = (
98+
ffmpeg
99+
.input('pipe:', format='rawvideo', pix_fmt='rgb24', s='{}x{}'.format(frame_width, frame_height))
100+
# -vf "drawtext=fontfile=Arial.ttf: fontsize=48: text=%{n}: x=(w-tw)/2: y=h-(2*lh): fontcolor=white: box=1: boxcolor=0x00000099"
101+
.drawtext(text="%{n}", escape_text=False,
102+
#x=50, y=50,
103+
x="(w-tw)/2", y="h-(2*lh)",
104+
fontfile="Arial.ttf", fontsize=font_width, fontcolor="white",
105+
#boxcolor="0x00000099",
106+
box=1, boxborderw=2, boxcolor="[email protected]")
107+
108+
.output(str(outfile), pix_fmt='yuv420p')
109+
110+
.overwrite_output()
111+
.run_async(pipe_stdin=True)
112+
)
113+
114+
assert frame_width is not None and frame_height is not None and ffmpeg_video_out_process is not None
115+
116+
# Send the frame to the ffmpeg process
117+
ffmpeg_video_out_process.stdin.write(img.tobytes())
118+
119+
elif gen == "Generated":
120+
121+
# The first frame can NOT be a generated one
122+
assert frame_width is not None and frame_height is not None
123+
124+
# Create an artificial black frame
125+
print(f"Injecting Black frame at idx {idx}")
126+
black_frame = np.zeros((frame_height, frame_width, 3), dtype=np.uint8)
127+
ffmpeg_video_out_process.stdin.write(black_frame.tobytes())
128+
129+
else:
130+
raise Exception(f"Unexpected value '{gen}' in column 'generated' at index {idx}")
131+
132+
# Close the video stream
133+
if ffmpeg_video_out_process is not None:
134+
ffmpeg_video_out_process.stdin.close()
135+
ffmpeg_video_out_process.wait()
136+
ffmpeg_video_out_process = None

0 commit comments

Comments
 (0)