batch_processing.py

# Imports
import time
import matplotlib # type: ignore
import matplotlib.pyplot as plt # type: ignore
import numpy as np
import argparse
import pandas as pd # type: ignore
import os
import sys
import cv2 # type: ignore
import csv

def main():
    global num_chirps, num_samples, min_scale, max_scale, max_theoretical_vel, max_doppler_vel, cfar_params, dist, velocity_bins, optimal_max_scale, optimal_min_scale, imported_max_scale, imported_min_scale, R_res, v_res
    # Parse the input file argument
    # parser = argparse.ArgumentParser(description='Process FMCW Range Data')
    # parser.add_argument('input_file', type=str, help='Path to the input .npy file')
    # args = parser.parse_args()

    f = "saved_radar_data.npy"
    config = np.load(f[:-4]+"_config.npy")          # these files are generated by the "Range_Doppler_Plot.py" program
    all_data = np.load(f)
    # base_dir = os.path.join(*f.split('/')[:-1])
    #change name every time you run, this creates a folder
    output_dir, no_filter_dir, MTI_dir, CFAR_dir, combined_dir = create_output_dirs("rangedop_squaretest")

    MTI_filters = ['none', '2pulse', '3pulse']
    cfar_methods = ['none', 'average', 'greatest', 'smallest']

    # CFAR parameters
    cfar_params = {
                'average': {'num_guard_cells_range': 3, 'num_guard_cells_doppler': 3, 'num_ref_cells_range': 5, 'num_ref_cells_doppler': 5, 'bias': 1, 'cfar_method': 'average'},
                'greatest': {'num_guard_cells_range': 1, 'num_guard_cells_doppler': 1, 'num_ref_cells_range': 3, 'num_ref_cells_doppler': 3, 'bias': .2, 'cfar_method': 'greatest'},
                'smallest': {'num_guard_cells_range': 2, 'num_guard_cells_doppler': 2, 'num_ref_cells_range': 4, 'num_ref_cells_doppler': 4, 'bias': 1.9, 'cfar_method': 'smallest'},
                }

    max_dist = 10
    custom_filter = False

    optimal_min_scale = 4.5
    optimal_max_scale = 10
    if custom_filter:
        min_scale = optimal_min_scale
        max_scale = optimal_max_scale


    sample_rate = config[0]
    signal_freq = config[1]
    output_freq = config[2]
    num_chirps = int(config[3])
    chirp_BW = config[4]
    ramp_time_s = config[5]
    frame_length_ms = config[6]
    max_doppler_vel = config[7] if len(config) > 7 else 2
    if not custom_filter:
        max_scale = config[8] if len(config) > 8 else 10 
        min_scale = config[9] if len(config) > 9 else 0 
        imported_max_scale = config[8] if len(config) > 8 else 10 
        imported_min_scale = config[9] if len(config) > 9 else 0
        # print(imported_min_scale, imported_max_scale)
        if min_scale > max_scale:
            min_scale = 0
            imported_min_scale = 0
    min_doppler_plot_vel = config[10] if len(config) > 10 else 2 
    string_length = config[11] if len(config) > 11 else 2 
    sample_goal = config[12] if len(config) > 12 else 2 
    dist_from_centroid = config[13] if len(config) > 13 else 2 
    time_data = pd.read_csv(f"{f[:-4]}_time.csv")
    num_samples = len(all_data[0][0])

    data_start_time = time_data.iloc[0, 0]  # Get first timestamp

    PRI = frame_length_ms / 1e3
    PRF = 1 / PRI

    # Split into frames
    N_frame = int(PRI * float(sample_rate))

    # Obtain range-FFT x-axis
    c = 3e8
    wavelength = c / output_freq
    slope = chirp_BW / ramp_time_s
    freq = np.linspace(-sample_rate/2, sample_rate/2, N_frame)
    dist = (freq - signal_freq - 25e3) * c / (2 * slope)
    velocity_bins = np.linspace(-max_doppler_vel, max_doppler_vel, num_chirps)

    # Resolutions
    R_res = c / (2 * chirp_BW)
    # print(f"Range resolution: {R_res} m")
    v_res = wavelength / (2 * num_chirps * PRI)
    # print(f"Velocity resolution: {v_res:.2f} m/s")

    # Doppler spectrum limits
    max_doppler_freq = PRF / 2
    # max_doppler_vel = max_doppler_freq * wavelength / 2

    # print("sample_rate = ", sample_rate/1e6, "MHz, ramp_time = ", int(ramp_time_s*(1e6)), "us, num_chirps = ", num_chirps, ", PRI = ", frame_length_ms, " ms")
    max_theoretical_vel = wavelength * PRF / 4
    # print(f"Theoretical max velocity: ±{max_theoretical_vel:.2f} m/s")

    # Plot range doppler data, loop through at the end of the data set
    cmn = 'viridis'
    
    process_filter('none', 'none', all_data, no_filter_dir)

    for MTI_opt in MTI_filters:
        if not MTI_opt == 'none':
            process_filter(MTI_opt, 'none', all_data, MTI_dir)
    for cfar_opt in cfar_methods:
        if not cfar_opt == 'none':
            process_filter('none', cfar_opt, all_data, CFAR_dir)

    for MTI_opt in MTI_filters:
        for cfar_opt in cfar_methods:
            if not MTI_opt == 'none' and not cfar_opt == 'none':
                process_filter(MTI_opt, cfar_opt, all_data, combined_dir)
    
def process_filter(MTI_opt, cfar_opt, all_data, output_dir):
    global num_chirps, num_samples, cfar_params, min_scale, max_scale, dist, optimal_min_scale, optimal_max_scale, imported_max_scale, imported_min_scale, R_res
    time_sum = 0

    if MTI_opt == 'none' and cfar_opt == 'none':
        csv_file_name = f"{output_dir}processed_data.csv"
        min_scale = optimal_min_scale
        max_scale = optimal_max_scale
    elif MTI_opt == 'none' and cfar_opt != 'none':
        csv_file_name = f"{output_dir}/{cfar_opt}/processed_data.csv"
        min_scale = imported_min_scale
        max_scale = imported_max_scale
    elif cfar_opt == 'none' and MTI_opt != 'none':
        csv_file_name = f"{output_dir}/{MTI_opt}/processed_data.csv"
        min_scale = optimal_min_scale
        max_scale = optimal_max_scale
    else:
        csv_file_name = f"{output_dir}{cfar_opt}_{MTI_opt}/processed_data.csv"
        min_scale = imported_min_scale
        max_scale = imported_max_scale

    
    with open(csv_file_name, 'w',newline='') as file:
        writer = csv.writer(file)
        writer.writerow(['Frame', 'Processing Time (s)', 'Total Time (s)', 'Peak Range (m)', 'Peak Velocity (m/s)', 'Peak Magnitude'])

    start_time = time.time()
    for i in range(len(all_data)):
        start_proc_time = time.time()
        radar_data = all_data[i]
        range_doppler_data = freq_process(radar_data)

        if MTI_opt == 'none':
            mti_filtered_data = range_doppler_data
        else:
            Chirp2P, Chirp3P = pulse_canceller(radar_data)
            if MTI_opt == '2pulse':
                mti_filtered_data = freq_process(Chirp2P)
            elif MTI_opt == '3pulse':
                mti_filtered_data = freq_process(Chirp3P)

        if cfar_opt == 'none':
            cfar_filtered_data = mti_filtered_data
        else:
            if cfar_opt == 'average':
                params = cfar_params[cfar_opt]
                # print(min_scale)
                CFAR_avg_parameters = {**params, 'min_value': min_scale}
                cfar_filtered_data, _ = cfar_2d(mti_filtered_data, **CFAR_avg_parameters)
            elif cfar_opt == 'greatest':
                params = cfar_params[cfar_opt]
                CFAR_greatest_parameters = {**params, 'min_value': min_scale}
                cfar_filtered_data, _ = cfar_2d(mti_filtered_data, **CFAR_greatest_parameters)
            elif cfar_opt == 'smallest':
                params = cfar_params[cfar_opt]
                CFAR_smallest_parameters = {**params, 'min_value': min_scale}
                cfar_filtered_data, _ = cfar_2d(mti_filtered_data, **CFAR_smallest_parameters)

        end_proc_time = time.time()
        output_file_title = f"{i:04d}.npy"
        img_file_title = f"{i:04d}.png"
        if MTI_opt == 'none' and cfar_opt == 'none':
            data_output_file_name = f"{output_dir}/Data/{output_file_title}"
            img_output_file_name = f"{output_dir}/Frames/{img_file_title}"
            # csv_file_name = f"{output_dir}processed_data.csv"
        elif MTI_opt == 'none' and cfar_opt != 'none':
            data_output_file_name = f"{output_dir}/{cfar_opt}/Data/{output_file_title}"
            img_output_file_name = f"{output_dir}/{cfar_opt}/Frames/{img_file_title}"
            # csv_file_name = f"{output_dir}/{cfar_opt}/processed_data.csv"
        elif cfar_opt == 'none' and MTI_opt != 'none':
            data_output_file_name = f"{output_dir}/{MTI_opt}/Data/{output_file_title}"
            img_output_file_name = f"{output_dir}/{MTI_opt}/Frames/{img_file_title}"
            # csv_file_name = f"{output_dir}/{MTI_opt}/processed_data.csv"
        else:
            data_output_file_name = f"{output_dir}{cfar_opt}_{MTI_opt}/Data/{output_file_title}"
            img_output_file_name = f"{output_dir}{cfar_opt}_{MTI_opt}/Frames/{img_file_title}"
            # csv_file_name = f"{output_dir}{cfar_opt}_{MTI_opt}/processed_data.csv"

        np.save(data_output_file_name, cfar_filtered_data)

        # Create png image
        nomalized_data = cv2.normalize(cfar_filtered_data, None, 0, 255, cv2.NORM_MINMAX)
        nomalized_data = np.uint8(nomalized_data)
        rows, cols = nomalized_data.shape
        vel_values = np.linspace(-max_theoretical_vel, max_theoretical_vel, cols)
        range_values = np.linspace(dist.min(), dist.max(), rows)
        vel_mask = np.abs(vel_values) <= max_doppler_vel
        start_col = np.argmax(vel_mask)
        end_col = cols - np.argmax(vel_mask[::-1])

        min_dist_for_56 = dist.max() - 56*R_res - R_res * 6
        # print(f"min_dist_for_56: {min_dist_for_56:.2f} m")
        # print(f"dist.max(): {dist.max():.2f} m")

        min_range_idx = np.argmin(np.abs(range_values - min_dist_for_56)) 
        cropped_data = nomalized_data[min_range_idx:, start_col:end_col]
        img = cv2.applyColorMap(cropped_data, cv2.COLORMAP_VIRIDIS)
        #cv2.imwrite(img_output_file_name, img)

        #changed to below to make processed images square

        # produce color image as before
        # ensure 3-channel for consistent handling
        if len(img.shape) == 2:
            img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)

        h, w = img.shape[:2]
        size = max(h, w)

        # scale so the shorter side becomes 'size'
        scale = size / min(h, w)
        new_w = int(round(w * scale))
        new_h = int(round(h * scale))
        img_resized = cv2.resize(img, (new_w, new_h), interpolation=cv2.INTER_AREA)

        # center crop to size x size
        start_x = (new_w - size) // 2
        start_y = (new_h - size) // 2
        img_square = img_resized[start_y:start_y + size, start_x:start_x + size]

        cv2.imwrite(img_output_file_name, img_square)



        proc_time = end_proc_time - start_proc_time
        time_sum += proc_time

        # print(max_doppler_vel)
        peak_range, peak_velocity, peak_magnitude = find_peak(cfar_filtered_data,range_values,vel_values, 
            max_doppler_vel=max_doppler_vel*.7,
            min_range=0,
            max_range=dist.max()
        )
        # print(max_doppler_vel)
        end_write_time = time.time()
        tot_time = end_write_time - start_time
        with open(csv_file_name, 'a',newline='') as file:
            writer = csv.writer(file)
            writer.writerow([i, proc_time, tot_time, peak_range, peak_velocity, peak_magnitude])

    print(f"Processed {len(all_data)} frames in {time_sum:.1f} s with MTI: {MTI_opt} and CFAR: {cfar_opt}")

def find_peak(radar_data, dist, vel_values, max_doppler_vel=None, min_range=-1, max_range=None):
    # Set default constraints if not provided
    if max_doppler_vel is None:
        max_doppler_vel = np.max(np.abs(vel_values))
    if max_range is None:
        max_range = np.max(dist)
        
    rows, cols = radar_data.shape
    
    # Find indices that satisfy the constraints
    range_indices = np.where((dist >= min_range) & (dist <= max_range))[0]
    vel_indices = np.where(np.abs(vel_values) <= max_doppler_vel)[0]
    
    if len(range_indices) == 0 or len(vel_indices) == 0:
        return None, None, None
    
    # Extract the constrained data
    constrained_data = radar_data[np.ix_(range_indices, vel_indices)]
    
    # Find the position of the maximum value
    if constrained_data.size > 0:
        max_idx = np.unravel_index(np.argmax(constrained_data), constrained_data.shape)
        peak_range_idx = range_indices[max_idx[0]]
        peak_vel_idx = vel_indices[max_idx[1]]
        
        peak_range = dist[peak_range_idx]
        peak_velocity = vel_values[peak_vel_idx]
        peak_magnitude = radar_data[peak_range_idx, peak_vel_idx]
        if peak_magnitude < 1:
            return None, None, None
        
        return peak_range, peak_velocity, peak_magnitude
    else:
        return None, None, None
    
def pulse_canceller(radar_data):
    global num_chirps, num_samples
    rx_chirps = []
    rx_chirps = radar_data
    # create 2 pulse canceller MTI array
    Chirp2P = np.empty([num_chirps, num_samples])*1j
    for chirp in range(num_chirps-1):
        chirpI = rx_chirps[chirp,:]
        chirpI1 = rx_chirps[chirp+1,:]
        chirp_correlation = np.correlate(chirpI, chirpI1, 'valid')
        angle_diff = np.angle(chirp_correlation, deg=False)  # returns radians
        Chirp2P[chirp,:] = (chirpI1 - chirpI * np.exp(-1j*angle_diff[0]))
    # create 3 pulse canceller MTI array
    Chirp3P = np.empty([num_chirps, num_samples])*1j
    for chirp in range(num_chirps-2):
        chirpI = Chirp2P[chirp,:]
        chirpI1 = Chirp2P[chirp+1,:]
        Chirp3P[chirp,:] = chirpI1 - chirpI
    return Chirp2P, Chirp3P

def cfar_2d(range_doppler_data, num_guard_cells_range=4, num_guard_cells_doppler=2, 
           num_ref_cells_range=8, num_ref_cells_doppler=4, bias=3, cfar_method='average', 
           min_value=-200):
    range_size, doppler_size = range_doppler_data.shape
    filtered_data = np.full_like(range_doppler_data, min_value)
    threshold = np.zeros_like(range_doppler_data)
    
    # Define window sizes for processing
    guard_window_range = 2 * num_guard_cells_range + 1
    guard_window_doppler = 2 * num_guard_cells_doppler + 1
    ref_window_range = guard_window_range + 2 * num_ref_cells_range
    ref_window_doppler = guard_window_doppler + 2 * num_ref_cells_doppler
    
    # Calculate the number of reference cells
    num_ref_cells_total = ref_window_range * ref_window_doppler - guard_window_range * guard_window_doppler
    
    # For each cell under test (CUT)
    for i in range(range_size):
        for j in range(doppler_size):
            # Define window boundaries
            range_min = max(0, i - num_guard_cells_range - num_ref_cells_range)
            range_max = min(range_size, i + num_guard_cells_range + num_ref_cells_range + 1)
            doppler_min = max(0, j - num_guard_cells_doppler - num_ref_cells_doppler)
            doppler_max = min(doppler_size, j + num_guard_cells_doppler + num_ref_cells_doppler + 1)
            
            # Create mask for guard cells
            guard_range_min = max(0, i - num_guard_cells_range)
            guard_range_max = min(range_size, i + num_guard_cells_range + 1)
            guard_doppler_min = max(0, j - num_guard_cells_doppler)
            guard_doppler_max = min(doppler_size, j + num_guard_cells_doppler + 1)
            
            # Extract window and create guards mask
            window = range_doppler_data[range_min:range_max, doppler_min:doppler_max]
            mask = np.ones_like(window, dtype=bool)
            
            # Mask out guard cells and CUT
            local_guard_range_min = max(0, guard_range_min - range_min)
            local_guard_range_max = local_guard_range_min + (guard_range_max - guard_range_min)
            local_guard_doppler_min = max(0, guard_doppler_min - doppler_min)
            local_guard_doppler_max = local_guard_doppler_min + (guard_doppler_max - guard_doppler_min)
            
            mask[local_guard_range_min:local_guard_range_max, 
                 local_guard_doppler_min:local_guard_doppler_max] = False
            
            # Extract reference cells using mask
            ref_cells = window[mask]
            
            # Skip if no reference cells (edge case)
            if len(ref_cells) == 0:
                continue
                
            # Calculate threshold based on method
            if cfar_method == 'average':
                noise_level = np.mean(ref_cells)
            elif cfar_method == 'greatest':
                noise_level = np.max(ref_cells)
            elif cfar_method == 'smallest':
                noise_level = np.min(ref_cells)
            else:
                noise_level = np.mean(ref_cells)  # Default to average
            
            # Apply bias to get threshold
            cell_threshold = noise_level + bias
            threshold[i, j] = cell_threshold
            
            # Check if CUT is above threshold
            if range_doppler_data[i, j] > cell_threshold:
                filtered_data[i, j] = range_doppler_data[i, j]
    
    return filtered_data, threshold

def freq_process(data):
    global max_theoretical_vel, max_doppler_vel, min_scale, max_scale, v_res
    rx_chirps_fft = np.fft.fftshift(abs(np.fft.fft2(data)))
    range_doppler_data = np.log10(rx_chirps_fft).T
    
    # Calculate total number of velocity bins
    num_vel_bins = range_doppler_data.shape[1]
    
    # Map each bin to its true velocity
    vel_bin_values = np.linspace(-max_theoretical_vel, max_theoretical_vel, num_vel_bins)
    
    # Find which bins are within our display range
    valid_bins = np.abs(vel_bin_values) <= (max_doppler_vel + 3 * v_res + 5 * v_res)
    
    # Zero out all bins outside our display range
    for i in range(len(valid_bins)):
        if not valid_bins[i]:
            range_doppler_data[:, i] = min_scale
    
    # Apply other processing
    num_good = len(range_doppler_data[:,0])   
    center_delete = 0  # delete ground clutter velocity bins around 0 m/s
    if center_delete != 0:
        for g in range(center_delete):
            end_bin = int(num_chirps/2+center_delete/2)
            range_doppler_data[:,(end_bin-center_delete+g)] = np.zeros(num_good)
    range_delete = 0   # delete the zero range bins (these are Tx to Rx leakage)
    if range_delete != 0:
        for r in range(range_delete):
            start_bin = int(len(range_doppler_data)/2)
            range_doppler_data[start_bin+r, :] = np.zeros(num_chirps)
    range_doppler_data = np.clip(range_doppler_data, min_scale, max_scale)  # clip the data to control the max spectrogram scale
    return range_doppler_data

def create_output_dirs(base_dir):
    frames_dir_name = "Frames"
    data_dir_name = "Data"

    output_dir = f"{base_dir}/ProcessedData"
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
    else:
        print(f"Output directory {output_dir} already exists. Files may be overwritten.")

    no_filter_dir = f"{output_dir}/No_Filter/"
    if not os.path.exists(no_filter_dir):
        os.makedirs(no_filter_dir, exist_ok=True)
        os.makedirs(f"{no_filter_dir}/{frames_dir_name}/", exist_ok=True)
        os.makedirs(f"{no_filter_dir}/{data_dir_name}/", exist_ok=True)

    MTI_dir = f"{output_dir}/MTI/"
    if not os.path.exists(MTI_dir):
        os.makedirs(MTI_dir, exist_ok=True)
        os.makedirs(f"{MTI_dir}/2pulse/", exist_ok=True)
        os.makedirs(f"{MTI_dir}/2pulse/{frames_dir_name}/", exist_ok=True)
        os.makedirs(f"{MTI_dir}/2pulse/{data_dir_name}/", exist_ok=True)
        os.makedirs(f"{MTI_dir}/3pulse/", exist_ok=True)
        os.makedirs(f"{MTI_dir}/3pulse/{frames_dir_name}/", exist_ok=True)
        os.makedirs(f"{MTI_dir}/3pulse/{data_dir_name}/", exist_ok=True)

    CFAR_dir = f"{output_dir}/CFAR/"
    cfar_methods = ['average', 'greatest', 'smallest']
    if not os.path.exists(CFAR_dir):
        os.makedirs(CFAR_dir, exist_ok=True)
        for method in cfar_methods:
            os.makedirs(f"{CFAR_dir}/{method}/", exist_ok=True)
            os.makedirs(f"{CFAR_dir}/{method}/{frames_dir_name}/", exist_ok=True)
            os.makedirs(f"{CFAR_dir}/{method}/{data_dir_name}/", exist_ok=True)

    mti_methods = ['2pulse', '3pulse']
    combined_dir = f"{output_dir}/Combined/"
    if not os.path.exists(combined_dir):
        os.makedirs(combined_dir, exist_ok=True)
        for method in mti_methods:
            for cfar_method in cfar_methods:
                os.makedirs(f"{output_dir}/Combined/{cfar_method}_{method}/", exist_ok=True)
                os.makedirs(f"{output_dir}/Combined/{cfar_method}_{method}/{frames_dir_name}/", exist_ok=True)
                os.makedirs(f"{output_dir}/Combined/{cfar_method}_{method}/{data_dir_name}/", exist_ok=True)
    
    return output_dir, no_filter_dir, MTI_dir, CFAR_dir, combined_dir

if __name__ == "__main__":
    main()