V-jinmgong/Emulation: Add implementation and evaluation scripts for RL-BWE emulation on AlphaRTC platform

.gitignore

@@ -5,4 +5,6 @@ log/
 workdir/
 *.yuv
 *.y4m
-*.wav
+*.wav
+origin_videos/
+distorted_videos/

emulation.sh

@@ -11,6 +11,10 @@ DEFAULT_AUTOCLOSE=120
 DEFAULT_VIDEO_HEIGHT=1080
 DEFAULT_VIDEO_WIDTH=1920
 DEFAULT_VIDEO_FPS=30
+DEFAULT_CROP_WIDTH=150
+DEFAULT_CROP_HEIGHT=60
+DEFAULT_CROP_X=900
+DEFAULT_CROP_Y=1020
 DEFAULT_IF_SAVE_MEDIA=true
 DEFAULT_RAMDISK_SIZE=24
 
@@ -33,33 +37,41 @@ if [ -z "$RECEIVER_PATTERNS_STR" ] || [ -z "$SENDER_PATTERNS_STR" ]; then
 fi
 
 # Parse input parameters
-while getopts "hm:v:a:t:c:W:H:f:s:l:d" opt; do
-  case $opt in
-    h) echo "Usage: `basename $0` [-m model] [-v video_file] [-a audio_file][-t trace_file] [-c autoclose] [-W video_width] [-H video_height] [-f video_fps] [-s if_save_media] [-l log_dir] [-d ramdisk_size_G]" >&2
-       exit 0 ;;
-    m) MODEL_DIR="$OPTARG" ;;
-    v) VIDEO_FILE="$OPTARG" ;;
-    a) AUDIO_FILE="$OPTARG" ;;
-    t) TRACE_FILE="$OPTARG" ;;
-    c) AUTOCLOSE="$OPTARG" ;;
-    W) VIDEO_WIDTH="$OPTARG" ;;
-    H) VIDEO_HEIGHT="$OPTARG" ;; 
-    f) VIDEO_FPS="$OPTARG" ;;
-    s) IF_SAVE_MEDIA="$OPTARG" ;;
-    l) LOG_DIR="$OPTARG" ;;
-    d) RAMDISK_SIZE="$OPTARG" ;;
-    \?) echo "Invalid option: -$OPTARG" >&2
-        exit 1 ;;
-    :) echo "Option -$OPTARG requires an argument." >&2
-       exit 1 ;;
-  esac
+while getopts "hm:v:a:t:c:W:H:f:s:l:d:x:y:w:h:" opt; do
+    case $opt in
+        h) echo "Usage: `basename $0` [-m model] [-v video_file] [-a audio_file][-t trace_file] [-c autoclose] [-W video_width] [-H video_height] [-f video_fps] [-s if_save_media] [-l log_dir] [-d ramdisk_size_G] [-x crop_x] [-y crop_y] [-w crop_width] [-h crop_height]" >&2
+             exit 0 ;;
+        m) MODEL_DIR="$OPTARG" ;;
+        v) VIDEO_FILE="$OPTARG" ;;
+        a) AUDIO_FILE="$OPTARG" ;;
+        t) TRACE_FILE="$OPTARG" ;;
+        c) AUTOCLOSE="$OPTARG" ;;
+        W) VIDEO_WIDTH="$OPTARG" ;;
+        H) VIDEO_HEIGHT="$OPTARG" ;; 
+        f) VIDEO_FPS="$OPTARG" ;;
+        s) IF_SAVE_MEDIA="$OPTARG" ;;
+        l) LOG_DIR="$OPTARG" ;;
+        d) RAMDISK_SIZE="$OPTARG" ;;
+        x) CROP_X="$OPTARG" ;;
+        y) CROP_Y="$OPTARG" ;;
+        w) CROP_WIDTH="$OPTARG" ;;
+        h) CROP_HEIGHT="$OPTARG" ;;
+        \?) echo "Invalid option: -$OPTARG" >&2
+                exit 1 ;;
+        :) echo "Option -$OPTARG requires an argument." >&2
+             exit 1 ;;
+    esac
 done
 
 # Set default values
 AUTOCLOSE=${AUTOCLOSE:-$DEFAULT_AUTOCLOSE}
 VIDEO_WIDTH=${VIDEO_WIDTH:-$DEFAULT_VIDEO_WIDTH}
 VIDEO_HEIGHT=${VIDEO_HEIGHT:-$DEFAULT_VIDEO_HEIGHT}
 VIDEO_FPS=${VIDEO_FPS:-$DEFAULT_VIDEO_FPS}
+CROP_WIDTH=${CROP_WIDTH:-$DEFAULT_CROP_WIDTH}
+CROP_HEIGHT=${CROP_HEIGHT:-$DEFAULT_CROP_HEIGHT}
+CROP_X=${CROP_X:-$DEFAULT_CROP_X}
+CROP_Y=${CROP_Y:-$DEFAULT_CROP_Y}
 IF_SAVE_MEDIA=${IF_SAVE_MEDIA:-$DEFAULT_IF_SAVE_MEDIA}
 RAMDISK_SIZE=${RAMDISK_SIZE:-$DEFAULT_RAMDISK_SIZE}
 
@@ -102,34 +114,45 @@ generate_media_config() {
 
 # Calculate score for video and network
 calculate_score() {
-    local receiver_log=$1
-    local src_video=$2
-    local dst_video=$3
-    local output_file=$4
-
-    python3 ${SCRIPT_DIR}/metrics/eval_video.py \
-            --src_video "$src_video" \
-            --dst_video "$dst_video" \
-            --output "$output_file" \
-            --frame_align_method "ocr" \
-            --video_size "${VIDEO_WIDTH}x${VIDEO_HEIGHT}" \
-            --pixel_format "420" \
-            --bitdepth "8" \
-            --fps "$VIDEO_FPS"
+    local src_video=$1
+    local dst_video=$2
+    local output_file=$3
 
-    python3 ${SCRIPT_DIR}/metrics/eval_network.py \
-            --dst_network_log "$receiver_log" \
-            --output "$output_file"
+    # Clear or create origin_video_dir and distorted_video_dir
+    mkdir -p "$SCRIPT_DIR/origin_videos"
+    mkdir -p "$SCRIPT_DIR/distorted_videos"
+    while [ "$(ls -A "$SCRIPT_DIR/origin_videos")" ]; do
+        rm -rf "$SCRIPT_DIR/origin_videos/*"
+    done
+    while [ "$(ls -A "$SCRIPT_DIR/distorted_videos")" ]; do
+        rm -rf "$SCRIPT_DIR/distorted_videos/*"
+    done
+    # make sure that the origin_video_dir and distorted_video_dir to be empty
+    python3 ${SCRIPT_DIR}/metrics/calc_scores.py \
+            --original_video_path "$src_video" \
+            --distorted_video_path "$dst_video" \
+            --origin_video_dir "$SCRIPT_DIR/origin_videos" \
+            --distorted_video_dir "$SCRIPT_DIR/distorted_videos" \
+            --output_json "$output_file" \
+            --video_width "$VIDEO_WIDTH" \
+            --video_height "$VIDEO_HEIGHT" \
+            --video_fps "$VIDEO_FPS" \
+            --crop_width "$CROP_WIDTH" \
+            --crop_height "$CROP_HEIGHT" \
+            --crop_x "$CROP_X" \
+            --crop_y "$CROP_Y" \
+            --israwvideo \
+            --supplement 
 }
 
-# Compress logs according to patterns
-compress_logs() {
+# filter logs according to patterns
+filter_logs() {
     local receiver_log=$1
     local sender_log=$2
     local base_input_dir=$3
     local base_output_dir=$4
 
-    python3 ${SCRIPT_DIR}/compress_logs.py \
+    python3 ${SCRIPT_DIR}/filter_logs.py \
             --receiver_log "$receiver_log" \
             --sender_log "$sender_log" \
             --receiver_patterns "$RECEIVER_PATTERNS_STR" \
@@ -240,7 +263,7 @@ emulation(){
     echo "The network configuration process $network_pid has been terminated"
     sleep 10
 
-    # Post-process logs: calculate score and compress logs
+    # Post-process logs: calculate score and filter logs
     printf "Calculating score\n"
     model_name=${model_name:-gcc}
     store_dir="$LOG_DIR/$base_name/$trace_name/$model_name"
@@ -251,12 +274,11 @@ emulation(){
     else
         touch "$store_dir/score.json"
     fi
-    calculate_score "$receiver_logging" "$workdir_video" "$save_video" "$store_dir/score.json"
+    calculate_score "$workdir_video" "$save_video" "$store_dir/score.json"
     printf "Score calculated\n"
-    printf "Compressing logs\n"
-    compress_logs "$receiver_logging" "$sender_logging" "$trace_dir" "$store_dir"
-    printf "Logs compressed and moved\n"
-
+    printf "Filtering logs\n"
+    filter_logs "$receiver_logging" "$sender_logging" "$trace_dir" "$store_dir"
+    printf "Logs filtered and moved\n"
     # Clean up
     rm -r "$trace_dir"/*
     printf "Cleared transmission outputs\n"
@@ -281,8 +303,12 @@ main() {
     echo "VIDEO_WIDTH: $VIDEO_WIDTH"
     echo "VIDEO_HEIGHT: $VIDEO_HEIGHT"
     echo "VIDEO_FPS: $VIDEO_FPS"
-    echo "IF_SAVE_MEDIA: $IF_SAVE_MEDIA"
     echo "LOG_DIR: $LOG_DIR"
+    echo "CROP_X: $CROP_X"
+    echo "CROP_Y: $CROP_Y"
+    echo "CROP_WIDTH: $CROP_WIDTH"
+    echo "CROP_HEIGHT: $CROP_HEIGHT"
+    echo "IF_SAVE_MEDIA: $IF_SAVE_MEDIA"
     echo "RAMDISK_SIZE: $RAMDISK_SIZE"
     emulation
     # Add the rest of your script logic here

compress_logs.py → filter_logs.py

@@ -3,7 +3,7 @@
 import argparse
 import os
 
-def compress_logs(log_file, patterns, output_path):
+def filter_logs(log_file, patterns, output_path):
     patterns = [re.compile(pattern) for pattern in patterns]
     if os.path.exists(output_path):
         os.remove(output_path)
@@ -21,7 +21,7 @@ def get_output_path(input_path, base_input_dir, base_output_dir):
     return os.path.join(base_output_dir, relative_path)
 
 if __name__ == "__main__":
-    parser = argparse.ArgumentParser(description='Process and compress log files.')
+    parser = argparse.ArgumentParser(description='Process and filter log files.')
     parser.add_argument('--receiver_log', type=str, required=True, help='receiver log file path')
     parser.add_argument('--sender_log', type=str, required=True, help='sender log file path')
     parser.add_argument('--receiver_patterns', type=lambda s: [p for p in s.split('|') if p], help='pipe-separated list of receiver log regex patterns')
@@ -40,8 +40,8 @@ def get_output_path(input_path, base_input_dir, base_output_dir):
 
     output_path = get_output_path(receiver_log, base_input_dir, base_output_dir)
     os.makedirs(os.path.dirname(output_path), exist_ok=True)
-    compress_logs(receiver_log, receiver_patterns, output_path)
+    filter_logs(receiver_log, receiver_patterns, output_path)
 
     output_path = get_output_path(sender_log, base_input_dir, base_output_dir)
     os.makedirs(os.path.dirname(output_path), exist_ok=True)
-    compress_logs(sender_log, sender_patterns, output_path)
+    filter_logs(sender_log, sender_patterns, output_path)

metrics/calc_scores.py

@@ -0,0 +1,115 @@
+import os 
+import re
+import json
+import argparse
+import subprocess
+from utils.preprocess import *
+
+def calc_vmaf(origin_video, distorted_video, output_xml="vmaf.xml"):
+    cmd_result = subprocess.run([
+        "vmaf", "--reference", origin_video, "--distorted", distorted_video, "--output", output_xml
+    ], capture_output=True, text=True)
+    with open(output_xml, "r") as f:
+        re_result = re.search(r'metric name="vmaf".*?mean="([\d]+\.[\d]+)"', f.read())
+        if not re_result:
+            raise ValueError("Can not get vmaf score from terminal output")
+        vmaf_score = float(re_result.group(1))
+    return vmaf_score
+
+def calculate_psnr_ssim(origin_video, distorted_video):
+    # PSNR
+    result = subprocess.run([
+        'ffmpeg', '-i', origin_video, '-i', distorted_video,
+        '-lavfi', 'psnr', '-f', 'null', '-'
+    ], capture_output=True, text=True)
+    psnr_values = re.findall(r'average:(\d+\.\d+|inf)', result.stderr)
+    if psnr_values and psnr_values[0] == 'inf':
+        psnr_values[0] = float('inf')
+    psnr_score = float(psnr_values[0]) if psnr_values else None
+    # SSIM
+    result = subprocess.run([
+        'ffmpeg', '-i', origin_video, '-i', distorted_video,
+        '-lavfi', 'ssim', '-f', 'null', '-'
+    ], capture_output=True, text=True)
+    ssim_values = re.findall(r'All:(\d+\.\d+)', result.stderr)
+    ssim_score = float(ssim_values[0]) if ssim_values else None
+
+    return psnr_score, ssim_score
+
+def calculate_frame_loss_rate(original_frame_dir, distorted_frame_dir):
+    original_frames = set(f for f in os.listdir(original_frame_dir) if f.endswith('.png'))
+    distorted_frames = set(f for f in os.listdir(distorted_frame_dir) if f.endswith('.png'))
+
+    lost_frames = original_frames - distorted_frames
+    total_frames = len(original_frames)
+    lost_frame_count = len(lost_frames)
+
+    if total_frames == 0:
+        return 0.0
+
+    loss_rate = (lost_frame_count / total_frames) * 100
+    return loss_rate
+
+def calculate_metrics(original_video_path, distorted_video_path, origin_video_dir, distorted_video_dir, output_json,
+                      video_width, video_height, video_fps, 
+                      crop_width, crop_height, crop_x, crop_y,
+                      israwvideo=True, supplement=True):
+
+    extract_frames(original_video_path, origin_video_dir, width=video_width, height=video_height, fps=video_fps, israwvideo=israwvideo)
+    extract_frames(distorted_video_path, distorted_video_dir,fps=video_fps)
+
+    frame_loss_rate = calculate_frame_loss_rate(origin_video_dir, distorted_video_dir)
+    print(f'Frame loss rate: {frame_loss_rate}%')
+
+    recognize_frame_numbers(origin_video_dir, crop_width, crop_height, crop_x, crop_y)
+    recognize_frame_numbers(distorted_video_dir, crop_width, crop_height, crop_x, crop_y)
+
+    if supplement:
+        supplement_align_frames(origin_video_dir, distorted_video_dir)
+    else:
+        retain_common_frames(origin_video_dir, distorted_video_dir)
+
+    merge_frames(origin_video_dir, distorted_video_dir, video_fps)
+    vmaf = calc_vmaf(f'{origin_video_dir}/original.y4m', f'{distorted_video_dir}/distorted.y4m')
+    psnr, ssim = calculate_psnr_ssim(f'{origin_video_dir}/original.y4m', f'{distorted_video_dir}/distorted.y4m')
+
+    print(f'VMAF score: {vmaf}')
+    print(f'PSNR: {psnr}')
+    print(f'SSIM: {ssim}')
+
+    results = {
+        "Frame loss rate": frame_loss_rate,
+        "PSNR": psnr,
+        "SSIM": ssim,
+        "Vmaf score": vmaf,
+    }
+
+    with open(output_json, 'w') as f:
+        json.dump(results, f)
+        f.write('\n')
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser(description='Calculate video quality scores.')
+    parser.add_argument('--original_video_path', type=str, required=True, help='Path to the original video.')
+    parser.add_argument('--distorted_video_path', type=str, required=True, help='Path to the distorted video.')
+    parser.add_argument('--origin_video_dir', type=str, required=True, help='Directory to save frames from the original video.')
+    parser.add_argument('--distorted_video_dir', type=str, required=True, help='Directory to save frames from the distorted video.')
+    parser.add_argument('--output_json', type=str, required=True, help='Path to the output JSON file.')
+    parser.add_argument('--video_width', type=int, required=True, help='Width of the video.')
+    parser.add_argument('--video_height', type=int, required=True, help='Height of the video.')
+    parser.add_argument('--video_fps', type=int, required=True, help='Frames per second of the video.')
+    parser.add_argument('--crop_width', type=int, default=150, help='Width of the frame number area.')
+    parser.add_argument('--crop_height', type=int, default=60, help='Height of the frame number area.')
+    parser.add_argument('--crop_x', type=int, default=900, help='X coordinate of the frame number area.')
+    parser.add_argument('--crop_y', type=int, default=1020, help='Y coordinate of the frame number area.')
+    parser.add_argument('--israwvideo', action='store_true', help='Flag to indicate if the original video is raw YUV.')
+    parser.add_argument('--supplement', action='store_true', help='Flag to indicate if the distorted video needs to be supplemented.')
+
+    args = parser.parse_args()
+
+    calculate_metrics(
+        args.original_video_path, args.distorted_video_path, args.origin_video_dir, 
+        args.distorted_video_dir, args.output_json, args.video_width, args.video_height, 
+        args.video_fps, args.crop_width, args.crop_height, args.crop_x, args.crop_y,
+        args.israwvideo, args.supplement
+    )