dataloader.py

import os
import numpy as np
from torch.utils.data import Dataset, DataLoader
import torch

def log_and_scale(ffts, vmin=None, vmax=None):
    log_ffts = ffts.reshape(-1,ffts.shape[2])
    log_ffts = np.log10(log_ffts + 1e-6)
    if vmin is None:
        vmin = np.percentile(log_ffts, 10)
    if vmax is None:
        vmax = np.percentile(log_ffts, 98)
    normalized_ffts = (log_ffts - vmin) / (vmax - vmin)
    normalized_ffts = np.clip(normalized_ffts, -1, 2)
    ffts = normalized_ffts.reshape(ffts.shape)
    return ffts, vmin, vmax

class SoundCamDataset(Dataset):
    def __init__(self, filenames, model="soundcam",chunk_size=4096, nperseg=512, noverlap=256, rate=44100, channels=16):
        self.filenames = filenames
        self.model = model
        self.rate = rate
        self.chunk_size = chunk_size
        self.nperseg = nperseg
        self.noverlap = noverlap
        self.channels = channels
        self.fc = 6000
        
        # Load data distribution file here
        self.data_distribution()
        
        self.create_indices()

    def __len__(self):
        return len(self.indices)
    
    def create_indices(self):
        def create_indices_for_event(touch, size):
            touch_folded = touch.unfold(0, size, 1).permute(0, 2, 1)
            touch_folded = touch_folded.sum(dim=1)[:,0]
            valid = (touch_folded == 0) | (touch_folded == size)
            start_idx = torch.where(valid)[0]
            end_idx = torch.where(valid)[0] + size
            idx_sample = torch.stack((start_idx, end_idx), dim=1)
            return idx_sample
        
        def create_indices_for_event_with_padding(touch, size):
            touch_folded = touch.unfold(0, size+2, 1).permute(0, 2, 1)
            touch_folded = touch_folded.sum(dim=1)[:,0]
            valid = (touch_folded == 0) | (touch_folded == size+2)
            start_idx = torch.where(valid)[0] + 1
            end_idx = start_idx + size
            idx_sample = torch.stack((start_idx, end_idx), dim=1)
            return idx_sample

        self.indices = []
        self.minmax_fft = []
        for idx_file, filename in enumerate(self.filenames):
            data = np.load(filename)
            touch = data['touch']
            touch = torch.tensor(touch)
            if 'tap' in filename:
                idx_sample = create_indices_for_event(touch, size=3)
                length = len(idx_sample)
            elif 'widget' in filename:
                idx_sample = create_indices_for_event(touch, size=6)
                length = len(idx_sample)
            elif 'object' in filename:
                idx_sample = create_indices_for_event(touch, size=6)
                length = len(idx_sample)
            elif 'micro' in filename:
                idx_sample = create_indices_for_event(touch, size=3)
                length = len(idx_sample)
            elif 'swipe' in filename:
                idx_sample = create_indices_for_event_with_padding(touch, size=3)
                length = len(idx_sample)
            else:
                length = len(touch)
                idx_sample = torch.arange(length)
                idx_sample = torch.stack((idx_sample, idx_sample+1), dim=1)
            idx_file = torch.full((length,), idx_file, dtype=torch.long).unsqueeze(1)
            idx = torch.cat((idx_file, idx_sample), dim=1)
            self.indices.append(idx)
            if "soundcam" in self.model:
                fft = data['fft'][:,:,70:]
            elif "baseline" in self.model:
                fft = data['fft_single'][:,:,70:]
            _, vmin, vmax = log_and_scale(fft)
            scaler = np.repeat(np.array([[vmin, vmax]]), length, axis=0)
            self.minmax_fft.append(scaler)
        if self.indices:
            self.indices = torch.cat(self.indices)
            self.minmax_fft = np.concatenate(self.minmax_fft, axis=0)
        return self.indices
    
    def data_distribution(self):
        self.data_dist = np.load(f'./best_models/data_distribution_{self.model.split("_")[-1]}.npz')
        self.fft_min = self.data_dist['fft_min']
        self.fft_max = self.data_dist['fft_max']
        self.fft_mean = self.data_dist['fft_mean']
        self.fft_std = self.data_dist['fft_std']
        self.heatmap_min = self.data_dist['heatmap_min']
        self.heatmap_max = self.data_dist['heatmap_max']
        self.heatmap_mean = self.data_dist['heatmap_mean']
        self.heatmap_std = self.data_dist['heatmap_std']
        self.fingertip_min = 0.00616749 # self.data_dist['fingertip_min']
        self.fingertip_max = 0.04235372 # self.data_dist['fingertip_max']
        self.fingertip_mean = 0.02306079 # self.data_dist['fingertip_mean']
        self.fingertip_std = 0.01834833 # self.data_dist['fingertip_std']
        
    def scaler(self, dat, input_type="fft", scaler_type="minmax", per_channel=True):
        def minmax_scaler(dat, _min, _max):
            dat_scaled = ((dat - _min) / (_max - _min))
            dat_scaled = np.clip(dat_scaled, -1, 2)
            return dat_scaled
        
        def standard_scaler(dat, _mean, _std):
            dat_scaled = ((dat - _mean) / _std)
            dat_scaled = np.clip(dat_scaled, -1, 1)
            return dat_scaled
        
        if per_channel:
            if scaler_type == "minmax" and input_type == "fft":
                _dat = dat.reshape(-1, dat.shape[2])
                return minmax_scaler(_dat, self.fft_min, self.fft_max).reshape(dat.shape)
            elif scaler_type == "minmax" and input_type == "heatmap":
                _dat = dat.reshape(-1, dat.shape[1] * dat.shape[2])
                return minmax_scaler(_dat, self.heatmap_min, self.heatmap_max).reshape(dat.shape)
            elif scaler_type == "minmax" and input_type == "fingertip":
                return minmax_scaler(dat, self.fingertip_min, self.fingertip_max).reshape(dat.shape)
            elif scaler_type == "standard" and input_type == "fft":
                _dat = dat.reshape(-1, dat.shape[2])
                return standard_scaler(_dat, self.fft_mean, self.fft_std).reshape(dat.shape)
            elif scaler_type == "standard" and input_type == "heatmap":
                _dat = dat.reshape(-1, dat.shape[1] * dat.shape[2])
                return standard_scaler(_dat, self.heatmap_mean, self.heatmap_std).reshape(dat.shape)
            elif scaler_type == "standard" and input_type == "fingertip":
                return standard_scaler(dat, self.fingertip_mean, self.fingertip_std).reshape(dat.shape)
        else:
            if scaler_type == "minmax" and input_type == "fft":
                _dat = dat.reshape(-1, dat.shape[2])
                return minmax_scaler(_dat, self.fft_min, self.fft_max).reshape(dat.shape)
            elif scaler_type == "minmax" and input_type == "heatmap":
                _dat = dat.reshape(dat.shape[0], -1)
                return minmax_scaler(_dat, self.heatmap_min, self.heatmap_max).reshape(dat.shape)
            elif scaler_type == "minmax" and input_type == "fingertip":
                return minmax_scaler(dat, self.fingertip_min, self.fingertip_max).reshape(dat.shape)
            elif scaler_type == "standard" and input_type == "fft":
                _dat = dat.reshape(-1, dat.shape[2])
                return standard_scaler(_dat, self.fft_mean, self.fft_std).reshape(dat.shape)
            elif scaler_type == "standard" and input_type == "heatmap":
                _dat = dat.reshape(-1, dat.shape[1] * dat.shape[2])
                return standard_scaler(_dat, self.heatmap_mean, self.heatmap_std).reshape(dat.shape)
            elif scaler_type == "standard" and input_type == "fingertip":
                return standard_scaler(dat, self.fingertip_mean, self.fingertip_std).reshape(dat.shape)

    def remap_touch_to_object_class(self, filename, touch):
        if "none" in filename:
            touch[:] = 0
        elif "drawing" in filename:
            touch[touch == 1] = 1
        elif "scrubbing" in filename:
            touch[touch == 1] = 2
        elif "sandpaper" in filename:
            touch[touch == 1] = 3
        elif "zipper" in filename:
            touch[touch == 1] = 4
        elif "scissor" in filename:
            touch[touch == 1] = 5
        elif "powerdrill" in filename:
            touch[touch == 1] = 6
        elif "e-toothbrush" in filename:
            touch[touch == 1] = 7
        elif "spray" in filename:
            touch[touch == 1] = 8
        return touch
    
    def __getitem__(self, idx):
        idx_file, idx_sample_start, idx_sample_end = self.indices[idx]
        filename = self.filenames[idx_file]
        data = np.load(filename)
        fft = data['fft']
        fft_single = data['fft_single']
        heatmap = data['heatmap']
        fingertip = data['fingertip']
        touch = data['touch']
        touch = np.around(touch, decimals=0).astype(int)

        fingertip = fingertip[:,20] # index fingertip
        # fingertip = fingertip[:,19] # thumb fingertip

        fingertip = np.diff(fingertip,axis=0)
        fingertip = np.linalg.norm(fingertip, axis=1).reshape(-1,1)
        fingertip = np.concatenate((np.zeros((1,1)), fingertip), axis=0)

        fft = fft[idx_sample_start:idx_sample_end]
        fft_single = fft_single[idx_sample_start:idx_sample_end]
        heatmap = heatmap[idx_sample_start:idx_sample_end]
        fingertip = fingertip[idx_sample_start:idx_sample_end]
        touch = touch[idx_sample_start:idx_sample_end][0]
        if "object" in filename:
            touch = self.remap_touch_to_object_class(filename, touch)

        fft = self.scaler(np.log10(fft + 1e-6), input_type="fft", scaler_type="minmax", per_channel=True)
        fft_single = self.scaler(np.log10(fft_single + 1e-6), input_type="fft", scaler_type="minmax", per_channel=True)
        heatmap = self.scaler(heatmap, input_type="heatmap", scaler_type="minmax", per_channel=True)
        fingertip = self.scaler(fingertip, input_type="fingertip", scaler_type="minmax", per_channel=False)
        fft = fft.reshape(1, -1, fft.shape[2])
        fft_single = fft_single.reshape(1, -1, fft_single.shape[2])
        heatmap = heatmap.reshape(1, -1, heatmap.shape[2])
        fingertip = fingertip.reshape(1, -1, fingertip.shape[1])

        fft = fft[:,:,70:]
        fft_single = fft_single[:,:,70:]

        fft = torch.tensor(fft, dtype=torch.float)
        fft_single = torch.tensor(fft_single, dtype=torch.float64).float()
        heatmap = torch.tensor(heatmap).float()
        fingertip = torch.tensor(fingertip).float()
        touch = torch.tensor(touch).float()

        if "soundcam" in self.model:
            return fft, heatmap, fingertip, touch
        elif "baseline" in self.model:
            return fft_single, fingertip, touch
        else:
            return fft, heatmap, fingertip, touch

def read_all_filenames(path):
    all_files = []
    for root, dirs, files in os.walk(path):
        for file in files:
            if file.endswith('.npz'):
                all_files.append(os.path.join(root, file))
    
    return all_files


if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument("--path", type=str, default="data_preprocessed/swipe/")
    args = parser.parse_args()

    filenames = read_all_filenames(args.path)

    import tqdm
    
    fft_min_max = []
    fft_mean_std = []
    heatmap_min_max = []
    heatmap_mean_std = []
    fingertip_min_max = []
    fingertip_mean_std = []
    for filename in tqdm.tqdm(filenames):
        data = np.load(filename)
        fft = data['fft'].reshape(-1, data['fft'].shape[2])
        fft = np.log10(fft + 1e-6)
        heatmap = data['heatmap'].reshape(data['heatmap'].shape[0], -1)
        fft_min_max.append([np.percentile(fft, 10, axis=0), np.percentile(fft, 98, axis=0)])
        fft_mean_std.append([np.mean(fft, axis=0), np.std(fft, axis=0)])
        heatmap_min_max.append([np.percentile(heatmap, 10, axis=0), np.percentile(heatmap, 90, axis=0)])
        heatmap_mean_std.append([np.mean(heatmap, axis=0), np.std(heatmap, axis=0)])
        fingertip = data['fingertip'][:,20]
        fingertip = np.diff(fingertip,axis=0)
        fingertip = np.linalg.norm(fingertip, axis=1)
        fingertip_min_max.append([np.percentile(fingertip, 10, axis=0), np.percentile(fingertip, 90, axis=0)])
        fingertip_mean_std.append([np.mean(fingertip, axis=0), np.std(fingertip, axis=0)])
    fft_min_max = np.array(fft_min_max)
    fft_mean_std = np.array(fft_mean_std)
    heatmap_min_max = np.array(heatmap_min_max)
    heatmap_mean_std = np.array(heatmap_mean_std)

    fft_min_max = np.mean(fft_min_max, axis=0)
    fft_mean_std = np.mean(fft_mean_std, axis=0)
    heatmap_min_max = np.mean(heatmap_min_max, axis=0)
    heatmap_mean_std = np.mean(heatmap_mean_std, axis=0)
    fingertip_min_max = np.mean(fingertip_min_max, axis=0)
    fingertip_mean_std = np.mean(fingertip_mean_std, axis=0)

    fft_min = fft_min_max[0]
    fft_max = fft_min_max[1]
    fft_mean = fft_mean_std[0]
    fft_std = fft_mean_std[1]
    heatmap_min = heatmap_min_max[0]
    heatmap_max = heatmap_min_max[1]
    heatmap_mean = heatmap_mean_std[0]
    heatmap_std = heatmap_mean_std[1]

    print(fft_min.shape, fft_max.shape, fft_mean.shape, fft_std.shape, heatmap_min.shape, heatmap_max.shape, heatmap_mean.shape, heatmap_std.shape)
    np.savez("./best_models/data_distribution.npz",
                fft_min=fft_min,
                fft_max=fft_max,
                fft_mean=fft_mean,
                fft_std=fft_std,
                heatmap_min=heatmap_min,
                heatmap_max=heatmap_max,
                heatmap_mean=heatmap_mean,
                heatmap_std=heatmap_std
            )