Generate automatically ECG-data with associated backgrounds.

ecg_generator.py

import matplotlib
import matplotlib.pyplot as plt
import numpy as np
import config as cfg
import neurokit2 as nk
import random
import os
import glob
import uuid

from PIL import Image

matplotlib.use("agg")


def check_path(path):
    paths = path.split("/")
    del paths[-1], paths[0]
    current_path = "../"

    for p in paths:
        current_path = current_path + "/" + str(p)
        if not os.path.exists(current_path):
            os.mkdir(current_path)


def plotter(active):
    def plot_result(func):
        if active:

            def wrapper():
                plt.plot(func()[0], func()[-1])
                plt.show()

            return wrapper
        else:
            return func

    return plot_result


def normalize_list(iterator, start, end):
    return [
        (
            (end - start) * (element - min(iterator)) / (max(iterator) - min(iterator))
            + start
        )
        for element in iterator
    ]


class HeartbeatBackgroundPlot:
    _parameters = cfg.parameters
    _image_parameters = cfg.image_parameters
    _hyper_parameters = cfg.hyper_parameters

    def __init__(
        self,
        heartbeat_simulation,
        number_samples,
        line_width,
        margin_x,
        margin_y,
        y_shift,
        origin,
        randint,
    ):
        self.heartbeat_simulation = heartbeat_simulation
        self.number_samples = number_samples
        self.line_width = line_width
        self.margin_x = margin_x
        self.margin_y = margin_y
        self.y_shift = y_shift
        self.origin = origin
        self.randint = randint

    def __set_color__(self, color):
        self.color = color

    def _save_image(
        self, background, x, y, parameters, color, label="", number=str(uuid.uuid4())
    ):
        fig, ax = plt.subplots()
        ax.axis("off")
        ax.set_frame_on(False)
        ax.imshow(background, **parameters)
        ax.plot(
            x,
            [e + self.y_shift for e in y][: self.number_samples],
            linewidth=self.line_width,
            color=color,
        )

        plt.style.use("dark_background")

        if label == "mask_":
            path = (
                self._parameters["data_path"] + "mask/ecg_artificial_" + label + number
            )
            check_path(path)
        else:
            path = (
                self._parameters["data_path"] + "img/ecg_artificial_" + label + number
            )
            check_path(path)

        fig.set_size_inches(self._parameters["inch"])
        plt.savefig(path, bbox_inches="tight", transparent=True, pad_inches=None)
        plt.close()
        plt.cla()
        plt.clf()

    def _save_image_mask(self, x, y, img):
        imshow_params = {
            "extent": [
                np.min(x) * self.margin_x,
                np.max(x) * self.margin_x,
                np.min(y) * self.margin_y,
                np.max(y) * self.margin_y,
            ],
            "aspect": "auto",
            "cmap": "binary",
            "interpolation": "nearest",
        }
        number = str(uuid.uuid4())

        self._save_image(
            background=img,
            x=x,
            y=y,
            parameters=imshow_params,
            label="origin_",
            color=self._parameters["color"],
            number=number,
        )
        self.__set_color__("white")
        self._save_image(
            background=Image.open(self._parameters["backgrounds"][0]),
            x=x,
            y=y,
            parameters=imshow_params,
            label="mask_",
            color="white",
            number=number,
        )

    def __brownian_motion(self, x):
        z = np.random.normal(self.origin, 1, self.number_samples)
        y = z[0] * x / (2 * np.pi) ** 0.5
        s = sum(np.sin(0.5 * n * y) * z[n] / n for n in range(1, self.number_samples))
        y += s * 2 / np.pi**0.5
        return y

    @plotter(cfg.plot_settings["decorators"])
    def random_brownian_heartbeat(self):
        # Theorem from Paley and Wiener says that a Fourier series
        # with random coefficients produces Brownian motion.
        x = np.linspace(self.origin, 2 * np.pi, self.number_samples)
        return [
            x,
            self.__brownian_motion(x),
        ]

    def create_imagepair_brownian_motion(self, image):
        img = Image.open(image)
        x, y = self.random_brownian_heartbeat()
        self._save_image_mask(x, y, img)

    def create_imagepair_heartbeat_background(self, image):
        img = Image.open(image)
        x = np.linspace(self.origin, 1, self.number_samples)

        if self.randint:
            y = nk.ecg_simulate(**self.heartbeat_simulation).tolist()
        else:
            y = nk.ecg_simulate(**self.heartbeat_simulation).tolist()

        self._save_image_mask(x=x, y=y, img=img)


if __name__ == "__main__":
    limit = int(input("Anzahl zu erzeugender Bilder: "))
    count = 0
    while len(glob.glob(cfg.parameters["data_mask"])) <= limit:
        try:
            heartbeatSimulation = dict(
                {
                    "duration": 10,
                    "noise": random.uniform(0.0, 0.4),
                    "heart_rate": int(random.uniform(10.0, 200.0)),
                    "method": "ecgsyn",
                }
            )

            plot = HeartbeatBackgroundPlot(
                heartbeat_simulation=heartbeatSimulation,
                number_samples=cfg.parameters["numberOfSamples"],
                origin=cfg.parameters["origin"],
                randint=cfg.parameters["random"],
                line_width=cfg.plot_settings["line_width"],
                margin_x=cfg.plot_settings["margin_x"],
                margin_y=cfg.plot_settings["margin_y"],
                y_shift=cfg.plot_settings["y_shift"],
            )
            j = random.randint(1, len(cfg.parameters["backgrounds"]) - 1)
            plot.create_imagepair_heartbeat_background(cfg.parameters["backgrounds"][j])
            count += 1

            if count % 10 == 0:
                plot.create_imagepair_brownian_motion(cfg.parameters["backgrounds"][j])

        except Exception as e:
            print("Raised has been: ", e)
            continue

This is the associated config.py file.

1. Put this file in the very same directory as the generator.
2. Update the corresponding background paths in parameters and/or add new ones.
3. Brownian motion will be generated about 10% of the sample number.

image_parameters = {
    "dpi": "figure",
    "bbox_inches": None,
    "pad_inches": 0.0,
    "facecolor": "black",
    "edgecolor": None,
    "backend": None,
}

parameters = {
    "origin": 0.0,
    "random_rate": 0.05,
    "rotation": (-0.1, 0.1),
    "backgrounds": [
        "./img/bg_black.png",
        "./img/bg_01.png",
        "./img/bg_02.png",
        "./img/bg_03.png",
        "./img/bg_04.png",
        "./img/bg_05.png",
    ],
    "shape": (256, 256, 1),
    "inch": (3.05, 3.065),
    "color": "gray",
}

This is the requirements.txt file:

matplotlib==3.5.1
neurokit2==0.1.7
numpy==1.22.3
Pillow==9.0.1
scipy==1.8.0