Skip to content

Commit ff31f02

Browse files
crnbakercrnbaker
committed
Addded scripts for colour reconstruction and denoising NanEye RAW files.
1 parent 0afb415 commit ff31f02

File tree

4 files changed

+357
-0
lines changed

4 files changed

+357
-0
lines changed

‎file_processing/create_matrix_video.py‎

Lines changed: 178 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,178 @@
1+
import os
2+
import subprocess
3+
import shutil
4+
5+
# --- Configuration ---
6+
INPUT_VIDEO = "../12th.avi"
7+
OUTPUT_VIDEO = "../settings_matrix.mp4"
8+
TEMP_RAW_DIR = "matrix_temp_raw"
9+
FRAME_WIDTH = 320
10+
FRAME_HEIGHT = 320
11+
12+
# --- Parameter Space Definition ---
13+
# Define the parameter space to scan.
14+
# The script will automatically scan any TWO parameters assigned a list of values.
15+
# The other two parameters should be assigned a single scalar value.
16+
17+
# Example: Scan Brightness and Contrast
18+
BRIGHTNESS = -0.15
19+
CONTRAST = 2.0
20+
SATURATION = [2.0, 2.0, 4.0, 6.0, 10.0] # Constant
21+
GAMMA = [0.6, 0.8, 1.0, 1.2, 1.4]
22+
23+
# Example: Scan Saturation and Gamma
24+
# BRIGHTNESS = 0.0 # Constant
25+
# CONTRAST = 1.5 # Constant
26+
# SATURATION = [0.5, 1.0, 1.5, 2.0, 2.5] # Scanned parameter (Y-axis)
27+
# GAMMA = [0.8, 0.9, 1.0, 1.1, 1.2] # Scanned parameter (X-axis)
28+
29+
30+
def extract_frames(video_path, output_dir):
31+
"""Extracts frames from the video and saves them as raw 16-bit files."""
32+
print(f"Extracting frames from {video_path}...")
33+
try:
34+
import cv2
35+
import numpy as np
36+
except ImportError as e:
37+
print(f"Import error: {e}. Please ensure opencv-python and numpy are installed.")
38+
return 0, 0
39+
40+
if os.path.exists(output_dir):
41+
shutil.rmtree(output_dir)
42+
os.makedirs(output_dir)
43+
44+
cap = cv2.VideoCapture(video_path)
45+
if not cap.isOpened():
46+
print(f"Error: Could not open video file {video_path}")
47+
return 0, 0
48+
49+
fps = cap.get(cv2.CAP_PROP_FPS)
50+
frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
51+
52+
for frame_num in range(frame_count):
53+
ret, frame = cap.read()
54+
if not ret:
55+
break
56+
if frame.shape[2] == 3:
57+
gray_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
58+
else:
59+
gray_frame = frame
60+
frame_16bit = np.uint16(gray_frame) * 256
61+
raw_frame_path = os.path.join(output_dir, f"frame_{frame_num:04d}.raw")
62+
frame_16bit.tofile(raw_frame_path)
63+
64+
cap.release()
65+
print(f"Extraction complete. {frame_count} frames extracted.")
66+
return frame_count, fps
67+
68+
69+
def create_matrix_video():
70+
if not os.path.exists(INPUT_VIDEO):
71+
print(f"Error: Input video not found at '{INPUT_VIDEO}'")
72+
return
73+
74+
# 1. Identify scanned and constant parameters
75+
all_params = {
76+
'brightness': BRIGHTNESS,
77+
'contrast': CONTRAST,
78+
'saturation': SATURATION,
79+
'gamma': GAMMA
80+
}
81+
scan_params = {k: v for k, v in all_params.items() if isinstance(v, (list, tuple))}
82+
const_params = {k: v for k, v in all_params.items() if not isinstance(v, (list, tuple))}
83+
84+
if len(scan_params) != 2:
85+
print(f"Error: Exactly two parameters must be specified as lists to be scanned.")
86+
print(f"Found {len(scan_params)} scanned parameters: {list(scan_params.keys())}")
87+
return
88+
89+
param_y_name, param_y_levels = list(scan_params.items())[0]
90+
param_x_name, param_x_levels = list(scan_params.items())[1]
91+
num_y = len(param_y_levels)
92+
num_x = len(param_x_levels)
93+
total_videos = num_y * num_x
94+
95+
# 2. Prepare raw frame data
96+
num_frames, fps = extract_frames(INPUT_VIDEO, TEMP_RAW_DIR)
97+
if num_frames == 0:
98+
return
99+
100+
# 3. Build the ffmpeg filter_complex string
101+
filter_chains = []
102+
stream_names = []
103+
104+
# Split the input stream into N copies
105+
split_filter = f"[0:v]split={total_videos}"
106+
for i in range(total_videos):
107+
stream_name = f"[s{i}]"
108+
split_filter += stream_name
109+
stream_names.append(stream_name)
110+
filter_chains.append(split_filter)
111+
112+
# Process each split stream
113+
processed_streams = []
114+
for i in range(total_videos):
115+
y_idx = i // num_x
116+
x_idx = i % num_x
117+
118+
current_params = const_params.copy()
119+
current_params[param_y_name] = param_y_levels[y_idx]
120+
current_params[param_x_name] = param_x_levels[x_idx]
121+
122+
in_stream = stream_names[i]
123+
out_stream = f"[v{i}]"
124+
processed_streams.append(out_stream)
125+
126+
eq_filter_parts = [f"{name}={value:.2f}" for name, value in current_params.items()]
127+
eq_filter = f"eq={':'.join(eq_filter_parts)}"
128+
129+
y_val = param_y_levels[y_idx]
130+
x_val = param_x_levels[x_idx]
131+
text_filter = f"drawtext=text='{param_y_name[0].upper()}={y_val:.2f} {param_x_name[0].upper()}={x_val:.2f}':x=10:y=10:fontsize=16:fontcolor=white"
132+
133+
filter_chains.append(f"{in_stream}{eq_filter},{text_filter}{out_stream}")
134+
135+
# Stack videos horizontally into rows
136+
row_streams = []
137+
for i in range(num_y):
138+
row_input_streams = processed_streams[i * num_x:(i + 1) * num_x]
139+
row_out_stream = f"[row{i}]"
140+
row_streams.append(row_out_stream)
141+
filter_chains.append(f"{''.join(row_input_streams)}hstack=inputs={num_x}{row_out_stream}")
142+
143+
# Stack rows vertically
144+
filter_chains.append(f"{''.join(row_streams)}vstack=inputs={num_y}[out]")
145+
146+
filter_complex_string = ";".join(filter_chains)
147+
148+
# 4. Build and run the final ffmpeg command
149+
ffmpeg_command = [
150+
"ffmpeg",
151+
"-framerate", str(fps),
152+
"-f", "image2",
153+
"-s", f"{FRAME_WIDTH}x{FRAME_HEIGHT}",
154+
"-pix_fmt", "bayer_bggr16le",
155+
"-i", os.path.join(TEMP_RAW_DIR, "frame_%04d.raw"),
156+
"-filter_complex", filter_complex_string,
157+
"-map", "[out]",
158+
"-c:v", "libx264",
159+
"-pix_fmt", "yuv420p",
160+
"-crf", "23",
161+
"-y",
162+
OUTPUT_VIDEO
163+
]
164+
165+
print("Running ffmpeg to create settings matrix video...")
166+
print(f"Scanning {param_y_name} (Y-axis) and {param_x_name} (X-axis).")
167+
print("This may take a while...")
168+
subprocess.run(ffmpeg_command)
169+
170+
# 5. Cleanup
171+
print("Cleaning up temporary files...")
172+
shutil.rmtree(TEMP_RAW_DIR)
173+
174+
print(f"Matrix video created successfully: {OUTPUT_VIDEO}")
175+
176+
177+
if __name__ == "__main__":
178+
create_matrix_video()

‎file_processing/denoise_video.py‎

Lines changed: 48 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,48 @@
1+
import cv2
2+
import numpy as np
3+
4+
5+
def denoise_video(input_video_path: str, output_video_path: str, kernel_size=3, side_by_side=False):
6+
7+
if kernel_size % 2 == 0:
8+
print("Error: kernel_size must be an odd number.")
9+
return
10+
11+
cap = cv2.VideoCapture(input_video_path)
12+
if not cap.isOpened():
13+
print(f"Error: Could not open video file {input_video_path}")
14+
return
15+
16+
frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
17+
frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
18+
fps = cap.get(cv2.CAP_PROP_FPS)
19+
20+
out_width = frame_width * 2 if side_by_side else frame_width
21+
22+
fourcc = cv2.VideoWriter_fourcc(*'mp4v')
23+
out = cv2.VideoWriter(output_video_path, fourcc, fps, (out_width, frame_height))
24+
25+
while True:
26+
ret, frame = cap.read()
27+
if not ret:
28+
break
29+
30+
denoised_frame = cv2.medianBlur(frame, kernel_size)
31+
32+
if side_by_side:
33+
combined_frame = np.concatenate((frame, denoised_frame), axis=1)
34+
out.write(combined_frame)
35+
else:
36+
out.write(denoised_frame)
37+
38+
cap.release()
39+
out.release()
40+
cv2.destroyAllWindows()
41+
42+
print(f"Processed video saved to {output_video_path}")
43+
44+
45+
if __name__ == "__main__":
46+
denoise_video("12th_color_reconstructed.mp4",
47+
"12th_color_reconstructed_denoised.mp4",
48+
kernel_size=9, side_by_side=True)

‎file_processing/reconstruct.py‎

Lines changed: 30 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,30 @@
1+
import os
2+
from reconstruction_utils import process_avi, process_awvideo
3+
4+
INPUT_FILE = "../12th.avi" # .avi or .awvideo
5+
OUTPUT_FILE = "../12th_color_reconstructed.mp4"
6+
7+
COLOR_SETTINGS = {
8+
"contrast": 3 ,
9+
"brightness": -0.05,
10+
"saturation": 4,
11+
"gamma_r": 1.2,
12+
"gamma_g": 1.2,
13+
"gamma_b": 1.2,
14+
}
15+
16+
17+
def main():
18+
_, extension = os.path.splitext(INPUT_FILE)
19+
20+
if extension.lower() == ".avi":
21+
process_avi(INPUT_FILE, OUTPUT_FILE, COLOR_SETTINGS)
22+
elif extension.lower() == ".awvideo":
23+
process_awvideo(INPUT_FILE, OUTPUT_FILE, COLOR_SETTINGS)
24+
print("Warning: .awvideo files recorded by NanEye Viewer appear to contain clipping.")
25+
else:
26+
print("Error: Extension must be .avi or .awvideo file.")
27+
28+
29+
if __name__ == "__main__":
30+
main()

‎file_processing/reconstruction_utils.py‎

Lines changed: 101 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,101 @@
1+
import os
2+
import subprocess
3+
import shutil
4+
import numpy as np
5+
import cv2
6+
7+
8+
def clear_directory(directory):
9+
if os.path.exists(directory):
10+
shutil.rmtree(directory)
11+
os.makedirs(directory)
12+
13+
14+
def reconstruct_colour(raw_frames_dir: str, frame_width_px: int, frame_height_px: int, fps: int, output_path:str,
15+
color_settings: dict) -> None:
16+
ffmpeg_command = [
17+
"ffmpeg",
18+
"-start_number", "0",
19+
"-f", "image2",
20+
"-s", f"{frame_width_px}x{frame_height_px}",
21+
"-pix_fmt", "bayer_bggr16le",
22+
"-framerate", str(fps),
23+
"-i", os.path.join(raw_frames_dir, "frame_%04d.raw"),
24+
"-vf", f"eq=contrast={color_settings['contrast']}:brightness={color_settings['brightness']}:"
25+
f"saturation={color_settings['saturation']}:gamma_r={color_settings['gamma_r']}:"
26+
f"gamma_g={color_settings['gamma_g']}:gamma_b={color_settings['gamma_b']},"
27+
f"atadenoise=0a=0.1:0b=0.1:1a=0.1:1b=0.1",
28+
"-c:v", "libx264",
29+
"-pix_fmt", "yuv420p",
30+
"-crf", "18",
31+
"-y",
32+
output_path
33+
]
34+
print("Running ffmpeg...")
35+
subprocess.run(ffmpeg_command, check=True)
36+
37+
38+
def process_avi(input_path: str, output_path: str, color_settings: dict) -> None:
39+
print(f"Processing AVI file: {input_path}")
40+
raw_frames_dir = "temp_raw_frames"
41+
clear_directory(raw_frames_dir)
42+
43+
cap = cv2.VideoCapture(input_path)
44+
if not cap.isOpened():
45+
print(f"Error: Could not open video file {input_path}")
46+
return
47+
48+
frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
49+
frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
50+
fps = cap.get(cv2.CAP_PROP_FPS)
51+
52+
frame_count = 0
53+
while True:
54+
ret, frame = cap.read()
55+
if not ret:
56+
break
57+
if frame.shape[2] == 3:
58+
gray_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
59+
else:
60+
gray_frame = frame
61+
frame_16bit = np.uint16(gray_frame) * 256
62+
raw_frame_path = os.path.join(raw_frames_dir, f"frame_{frame_count:04d}.raw")
63+
frame_16bit.tofile(raw_frame_path)
64+
frame_count += 1
65+
cap.release()
66+
67+
reconstruct_colour(raw_frames_dir, frame_width, frame_height, fps, output_path, color_settings)
68+
shutil.rmtree(raw_frames_dir)
69+
print(f"Finished processing. Output saved to {output_path}")
70+
71+
72+
def process_awvideo(input_path: str, output_path: str, color_settings: dict) -> None:
73+
"""
74+
Processes a .awvideo file.
75+
"""
76+
print(f"Processing AWVIDEO file: {input_path}")
77+
raw_frames_dir = "temp_raw_frames"
78+
clear_directory(raw_frames_dir)
79+
80+
frame_width = 320
81+
frame_height = 320
82+
bytes_per_pixel = 2
83+
frame_size = frame_width * frame_height * bytes_per_pixel
84+
header_size = 20
85+
fps = 10
86+
87+
with open(input_path, "rb") as f:
88+
f.seek(header_size)
89+
frame_count = 0
90+
while True:
91+
frame_data = f.read(frame_size)
92+
if not frame_data:
93+
break
94+
raw_frame_path = os.path.join(raw_frames_dir, f"frame_{frame_count:04d}.raw")
95+
with open(raw_frame_path, "wb") as raw_f:
96+
raw_f.write(frame_data)
97+
frame_count += 1
98+
99+
reconstruct_colour(raw_frames_dir, frame_width, frame_height, fps, output_path, color_settings)
100+
shutil.rmtree(raw_frames_dir)
101+
print(f"Finished processing. Output saved to {output_path}")

0 commit comments

Comments
 (0)