White balance

whitebalancer.py

#!/usr/bin/env python

from pathlib import Path
import os
import cv2
import argparse

# Supported OpenCV Backend Ubuntu 24.04.1 LTS
os.environ["QT_QPA_PLATFORM"] = "xcb"
os.environ["OPENCV_UI_BACKEND"] = "xcb"

# Passed arguments
parser = argparse.ArgumentParser()
parser.add_argument('-i', '--input', help='<Required> Input video', required=True)
parser.add_argument('-o', '--output_dir', help='Output directory', required=False, default='.')
parser.add_argument('-s', '--scaling', help='Scale applied to output for visualizing', required=False, default=0.5, type=float)
parser.add_argument('-r', '--fps', help='Output FPS', required=False, default=None, type=int)
args = parser.parse_args()

# Create Output directory
if not os.path.isdir(args.output_dir):
    os.makedirs(args.output_dir)
    print("Created output dir {}".format(args.output_dir))

cv2.namedWindow("Stream", cv2.WINDOW_NORMAL)
cap = cv2.VideoCapture(args.input)
frame_idx = 0
record = False
video_writer = None

# Get video properties
frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
fps = int(cap.get(cv2.CAP_PROP_FPS))

if args.fps != None:
    fps = args.fps 
    
# Loop to read frames from the video
while cap.isOpened():
    ret, frame = cap.read()
    if not ret:
        break

    # Create a SimpleWB object
    wb = cv2.xphoto.createSimpleWB()

    # Apply white balance
    balanced_frame = wb.balanceWhite(frame)

    if args.scaling != 1:
        show_frame = cv2.resize(
            balanced_frame, None, fx=args.scaling, fy=args.scaling, interpolation=cv2.INTER_LINEAR)
    else:
        show_frame = balanced_frame

    cv2.imshow("Stream", show_frame)

    key = cv2.waitKey(1) & 0xFF
    if key == ord("q"):
        break
    elif key == ord("s"):
        fname = os.path.join(args.output_dir, f"{Path(args.input).stem}_{frame_idx}.png")
        cv2.imwrite(fname, balanced_frame)
        print(f"Saved image to {fname}")
    elif key == ord("r"):
        if record:
            print("Stopping recording...")
            video_writer.release()
            video_writer = None
            record = False
        else:
            output_video_path = os.path.join(args.output_dir, f"{Path(args.input).stem}_recorded.avi")
            fourcc = cv2.VideoWriter_fourcc(*"XVID")
            video_writer = cv2.VideoWriter(output_video_path, fourcc, fps, (frame_width, frame_height))
            print(f"Recording started: {output_video_path}")
            record = True
    
    if record and video_writer is not None:
        video_writer.write(balanced_frame)

    frame_idx += 1

cap.release()
if video_writer is not None:
    video_writer.release()
cv2.destroyAllWindows()