ConstructImage.py

import os
import random
import argparse
import time
import sys
from datetime import timedelta
import numpy as np
import matplotlib

matplotlib.use("Agg")
import matplotlib.pyplot as plt
from PIL import Image
import torch

from utils.dataloader import IMUDataset_fft
from utils.util import change_windowsize


def set_global_seed(seed: int) -> None:
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False


def parse_args():
    p = argparse.ArgumentParser()
    p.add_argument("--data_dir", type=str, default="./dataset/4activity")
    p.add_argument("-d", "--dataset", type=str, default="SBHAR")
    p.add_argument("--data_file", type=str, default="data_20_120_5activity.npy")
    p.add_argument("--label_file", type=str, default="label_20_120_4activity.npy")
    p.add_argument("--L", type=int, default=120)
    p.add_argument("--seed", type=int, default=42)
    p.add_argument("--save_root", type=str, default="./imu_images_tf/4activity")
    p.add_argument("--log_every", type=int, default=50)
    p.add_argument("--quiet", action="store_true")
    return p.parse_args()


def build_color_mapping(num_channels: int):
    plt_colors = [
        "green","black","blue","brown","chartreuse","chocolate","coral","crimson",
        "blueviolet","darkblue","darkgreen","firebrick","gold","teal","grey","indigo",
        "steelblue","indianred","goldenrod","darkred","darkorange","magenta","maroon",
        "navy","olive","orange","orchid","pink","plum","purple","red","cornflowerblue",
        "sienna","darkkhaki","tan","dodgerblue","darkseagreen","cadetblue",
    ]
    if num_channels > len(plt_colors):
        plt_colors = plt_colors * (num_channels // len(plt_colors) + 1)
    return plt_colors[:num_channels]


def imu_to_image_time_freq(sample_id: int, imu_data: np.ndarray, save_dir: str) -> str:
    C, T = imu_data.shape
    colors = build_color_mapping(C)

    rows, cols = 6, 2
    cell_h, cell_w = 37, 112
    img_h = rows * cell_h
    img_w = cols * cell_w

    dpi = 100
    fig = plt.figure(figsize=(img_w / dpi, img_h / dpi), dpi=dpi)

    for i in range(C):
        sig = imu_data[i]
        color = colors[i]

        ax1 = fig.add_subplot(rows, cols, i * cols + 1)
        ax1.plot(np.arange(T), sig, color=color, linewidth=0.5, marker="*", markersize=1)
        ax1.set_xticks([])
        ax1.set_yticks([])

        fft_vals = np.fft.rfft(sig)
        fft_mag = np.abs(fft_vals)
        fft_mag = fft_mag / (fft_mag.max() + 1e-8)

        ax2 = fig.add_subplot(rows, cols, i * cols + 2)
        ax2.plot(fft_mag, color=color, linewidth=0.7)
        ax2.set_xticks([])
        ax2.set_yticks([])

    fig.subplots_adjust(top=1, bottom=0, left=0, right=1, wspace=0, hspace=0)

    os.makedirs(save_dir, exist_ok=True)
    save_path = os.path.join(save_dir, f"{sample_id}.png")
    fig.savefig(save_path, pad_inches=0)
    plt.close(fig)

    img = Image.open(save_path).convert("RGB")
    img.save(save_path)

    return save_path


def load_imu_whole_dataset(args):
    data_path = os.path.join(args.data_dir, args.dataset, args.data_file)
    label_path = os.path.join(args.data_dir, args.dataset, args.label_file)
    data = np.load(data_path).astype(np.float32)
    labels = np.load(label_path).astype(int)
    data, labels = change_windowsize(data, labels, args.L)
    return IMUDataset_fft(data, labels, isNormalization=True)


def extract_activity_label(y) -> int:
    if isinstance(y, torch.Tensor):
        return int(y.detach().cpu().numpy().reshape(-1)[0])
    return int(np.array(y).reshape(-1)[0])


def _format_td(seconds: float) -> str:
    if seconds < 0 or not np.isfinite(seconds):
        return "N/A"
    return str(timedelta(seconds=int(seconds)))


def convert_imu_to_images_all(args):
    set_global_seed(args.seed)
    dataset = load_imu_whole_dataset(args)

    save_dir = os.path.join(args.save_root, f"{args.dataset}_image")
    os.makedirs(save_dir, exist_ok=True)

    try:
        total = len(dataset)
    except TypeError:
        total = None

    all_labels = []
    sample_id = 0
    n_failed = 0
    t0 = time.time()

    if not args.quiet:
        print(f"[IMU2IMG] dataset={args.dataset} save_dir={save_dir} total={total if total is not None else 'unknown'}")

    for x_org, _, y in dataset:
        try:
            imu = x_org.numpy()
            imu_to_image_time_freq(sample_id, imu, save_dir)
            all_labels.append(extract_activity_label(y))
        except Exception as e:
            n_failed += 1
            if not args.quiet:
                print(f"[IMU2IMG][WARN] failed at sample_id={sample_id}: {repr(e)}", file=sys.stderr)
        finally:
            sample_id += 1

        if (not args.quiet) and args.log_every > 0 and sample_id % args.log_every == 0:
            elapsed = time.time() - t0
            speed = sample_id / max(elapsed, 1e-9)
            if total is not None and total > 0:
                pct = 100.0 * sample_id / total
                eta = (total - sample_id) / max(speed, 1e-9)
                print(
                    f"[IMU2IMG] {sample_id}/{total} ({pct:.2f}%) | "
                    f"speed={speed:.2f} samp/s | elapsed={_format_td(elapsed)} | "
                    f"eta={_format_td(eta)} | failed={n_failed}"
                )
            else:
                print(
                    f"[IMU2IMG] {sample_id} done | speed={speed:.2f} samp/s | "
                    f"elapsed={_format_td(elapsed)} | failed={n_failed}"
                )

    all_labels = np.asarray(all_labels, dtype=np.int64)
    np.save(os.path.join(save_dir, "labels.npy"), all_labels)

    if not args.quiet:
        elapsed = time.time() - t0
        print(f"[IMU2IMG] finished: saved={len(all_labels)} images, failed={n_failed}, elapsed={_format_td(elapsed)}")

    return save_dir, all_labels


def main():
    args = parse_args()
    convert_imu_to_images_all(args)


if __name__ == "__main__":
    main()