cwru_segmentation.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "cwru-segmentation.ipynb",
      "version": "0.3.2",
      "provenance": [],
      "toc_visible": true,
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/fboldt/SignalProcessing/blob/master/cwru_segmentation.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "O_WyaUYJoAAB"
      },
      "source": [
        "# CWRU files.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BxnouJISoB7q",
        "colab_type": "text"
      },
      "source": [
        "Associate each Matlab file name to a bearing condition in a Python dictionary. The dictionary keys identify the conditions.\n",
        "\n",
        "There are only four normal conditions, with loads of 0, 1, 2 and 3 hp. All conditions end with an underscore character followed by an algarism representing the load applied during the acquisitions. The remaining conditions follow the pattern:\n",
        "\n",
        "First two characters represent the bearing location, i.e. drive end (DE) and fan end (FE). The following two characters represent the failure location in the bearing, i.e. ball (BA), Inner Race (IR) and Outer Race (OR). The next three algarisms indicate the severity of the failure, where 007 stands for 0.007 inches and 0021 for 0.021 inches. For Outer Race failures, the character @ is followed by a number that indicates different load zones.\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ZVVoHvQEn8hi",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "debug = False\n",
        "# size of each segment\n",
        "sample_size = 32768\n",
        "if not debug:\n",
        "  sample_size = 512\n",
        "acquisitions = {}\n",
        "# Normal\n",
        "acquisitions[\"Normal_0\"] = \"97.mat\"\n",
        "acquisitions[\"Normal_1\"] = \"98.mat\"\n",
        "acquisitions[\"Normal_2\"] = \"99.mat\"\n",
        "acquisitions[\"Normal_3\"] = \"100.mat\"\n",
        "# DE Inner Race 0.007 inches\n",
        "acquisitions[\"DEIR.007_0\"] = \"105.mat\"\n",
        "acquisitions[\"DEIR.007_1\"] = \"106.mat\"\n",
        "acquisitions[\"DEIR.007_2\"] = \"107.mat\"\n",
        "acquisitions[\"DEIR.007_3\"] = \"108.mat\"\n",
        "# DE Ball 0.007 inches\n",
        "acquisitions[\"DEB.007_0\"] = \"118.mat\"\n",
        "acquisitions[\"DEB.007_1\"] = \"119.mat\"\n",
        "acquisitions[\"DEB.007_2\"] = \"120.mat\"\n",
        "acquisitions[\"DEB.007_3\"] = \"121.mat\"\n",
        "# DE Outer race 0.007 inches centered @6:00\n",
        "acquisitions[\"DEOR.007@6_0\"] = \"130.mat\"\n",
        "acquisitions[\"DEOR.007@6_1\"] = \"131.mat\"\n",
        "acquisitions[\"DEOR.007@6_2\"] = \"132.mat\"\n",
        "acquisitions[\"DEOR.007@6_3\"] = \"133.mat\"\n",
        "# DE Outer race 0.007 inches centered @3:00\n",
        "acquisitions[\"DEOR.007@3_0\"] = \"144.mat\"\n",
        "acquisitions[\"DEOR.007@3_1\"] = \"145.mat\"\n",
        "acquisitions[\"DEOR.007@3_2\"] = \"146.mat\"\n",
        "acquisitions[\"DEOR.007@3_3\"] = \"147.mat\"\n",
        "# DE Outer race 0.007 inches centered @12:00\n",
        "acquisitions[\"DEOR.007@12_0\"] = \"156.mat\"\n",
        "acquisitions[\"DEOR.007@12_1\"] = \"158.mat\"\n",
        "acquisitions[\"DEOR.007@12_2\"] = \"159.mat\"\n",
        "acquisitions[\"DEOR.007@12_3\"] = \"160.mat\"\n",
        "# DE Inner Race 0.014 inches\n",
        "acquisitions[\"DEIR.014_0\"] = \"169.mat\"\n",
        "acquisitions[\"DEIR.014_1\"] = \"170.mat\"\n",
        "acquisitions[\"DEIR.014_2\"] = \"171.mat\"\n",
        "acquisitions[\"DEIR.014_3\"] = \"172.mat\"\n",
        "# DE Ball 0.014 inches\n",
        "acquisitions[\"DEB.014_0\"] = \"185.mat\"\n",
        "acquisitions[\"DEB.014_1\"] = \"186.mat\"\n",
        "acquisitions[\"DEB.014_2\"] = \"187.mat\"\n",
        "acquisitions[\"DEB.014_3\"] = \"188.mat\"\n",
        "# DE Outer race 0.014 inches centered @6:00\n",
        "acquisitions[\"DEOR.014@6_0\"] = \"197.mat\"\n",
        "acquisitions[\"DEOR.014@6_1\"] = \"198.mat\"\n",
        "acquisitions[\"DEOR.014@6_2\"] = \"199.mat\"\n",
        "acquisitions[\"DEOR.014@6_3\"] = \"200.mat\"\n",
        "# DE Ball 0.021 inches\n",
        "acquisitions[\"DEB.021_0\"] = \"222.mat\"\n",
        "acquisitions[\"DEB.021_1\"] = \"223.mat\"\n",
        "acquisitions[\"DEB.021_2\"] = \"224.mat\"\n",
        "acquisitions[\"DEB.021_3\"] = \"225.mat\"\n",
        "# FE Inner Race 0.021 inches\n",
        "acquisitions[\"FEIR.021_0\"] = \"270.mat\"\n",
        "acquisitions[\"FEIR.021_1\"] = \"271.mat\"\n",
        "acquisitions[\"FEIR.021_2\"] = \"272.mat\"\n",
        "acquisitions[\"FEIR.021_3\"] = \"273.mat\"\n",
        "# FE Inner Race 0.014 inches\n",
        "acquisitions[\"FEIR.014_0\"] = \"274.mat\"\n",
        "acquisitions[\"FEIR.014_1\"] = \"275.mat\"\n",
        "acquisitions[\"FEIR.014_2\"] = \"276.mat\"\n",
        "acquisitions[\"FEIR.014_3\"] = \"277.mat\"\n",
        "# FE Ball 0.007 inches\n",
        "acquisitions[\"FEB.007_0\"] = \"282.mat\"\n",
        "acquisitions[\"FEB.007_1\"] = \"283.mat\"\n",
        "acquisitions[\"FEB.007_2\"] = \"284.mat\"\n",
        "acquisitions[\"FEB.007_3\"] = \"285.mat\"\n",
        "# DE Inner Race 0.021 inches\n",
        "acquisitions[\"DEIR.021_0\"] = \"209.mat\"\n",
        "acquisitions[\"DEIR.021_1\"] = \"210.mat\"\n",
        "acquisitions[\"DEIR.021_2\"] = \"211.mat\"\n",
        "acquisitions[\"DEIR.021_3\"] = \"212.mat\"\n",
        "# DE Outer race 0.021 inches centered @6:00\n",
        "acquisitions[\"DEOR.021@6_0\"] = \"234.mat\"\n",
        "acquisitions[\"DEOR.021@6_1\"] = \"235.mat\"\n",
        "acquisitions[\"DEOR.021@6_2\"] = \"236.mat\"\n",
        "acquisitions[\"DEOR.021@6_3\"] = \"237.mat\"\n",
        "# DE Outer race 0.021 inches centered @3:00\n",
        "acquisitions[\"DEOR.021@3_0\"] = \"246.mat\"\n",
        "acquisitions[\"DEOR.021@3_1\"] = \"247.mat\"\n",
        "acquisitions[\"DEOR.021@3_2\"] = \"248.mat\"\n",
        "acquisitions[\"DEOR.021@3_3\"] = \"249.mat\"\n",
        "# DE Outer race 0.021 inches centered @12:00\n",
        "acquisitions[\"DEOR.021@12_0\"] = \"258.mat\"\n",
        "acquisitions[\"DEOR.021@12_1\"] = \"259.mat\"\n",
        "acquisitions[\"DEOR.021@12_2\"] = \"260.mat\"\n",
        "acquisitions[\"DEOR.021@12_3\"] = \"261.mat\"\n",
        "# FE Inner Race 0.007 inches\n",
        "acquisitions[\"FEIR.007_0\"] = \"278.mat\"\n",
        "acquisitions[\"FEIR.007_1\"] = \"279.mat\"\n",
        "acquisitions[\"FEIR.007_2\"] = \"280.mat\"\n",
        "acquisitions[\"FEIR.007_3\"] = \"281.mat\"\n",
        "# FE Ball 0.014 inches\n",
        "acquisitions[\"FEB.014_0\"] = \"286.mat\"\n",
        "acquisitions[\"FEB.014_1\"] = \"287.mat\"\n",
        "acquisitions[\"FEB.014_2\"] = \"288.mat\"\n",
        "acquisitions[\"FEB.014_3\"] = \"289.mat\"\n",
        "# FE Ball 0.021 inches\n",
        "acquisitions[\"FEB.021_0\"] = \"290.mat\"\n",
        "acquisitions[\"FEB.021_1\"] = \"291.mat\"\n",
        "acquisitions[\"FEB.021_2\"] = \"292.mat\"\n",
        "acquisitions[\"FEB.021_3\"] = \"293.mat\"\n",
        "# FE Outer race 0.007 inches centered @6:00\n",
        "acquisitions[\"FEOR.007@6_0\"] = \"294.mat\"\n",
        "acquisitions[\"FEOR.007@6_1\"] = \"295.mat\"\n",
        "acquisitions[\"FEOR.007@6_2\"] = \"296.mat\"\n",
        "acquisitions[\"FEOR.007@6_3\"] = \"297.mat\"\n",
        "# FE Outer race 0.007 inches centered @3:00\n",
        "acquisitions[\"FEOR.007@3_0\"] = \"298.mat\"\n",
        "acquisitions[\"FEOR.007@3_1\"] = \"299.mat\"\n",
        "acquisitions[\"FEOR.007@3_2\"] = \"300.mat\"\n",
        "acquisitions[\"FEOR.007@3_3\"] = \"301.mat\"\n",
        "# FE Outer race 0.007 inches centered @12:00\n",
        "acquisitions[\"FEOR.007@12_0\"] = \"302.mat\"\n",
        "acquisitions[\"FEOR.007@12_1\"] = \"305.mat\"\n",
        "acquisitions[\"FEOR.007@12_2\"] = \"306.mat\"\n",
        "acquisitions[\"FEOR.007@12_3\"] = \"307.mat\"\n",
        "# FE Outer race 0.014 inches centered @3:00\n",
        "acquisitions[\"FEOR.014@3_0\"] = \"310.mat\"\n",
        "acquisitions[\"FEOR.014@3_1\"] = \"309.mat\"\n",
        "acquisitions[\"FEOR.014@3_2\"] = \"311.mat\"\n",
        "acquisitions[\"FEOR.014@3_3\"] = \"312.mat\"\n",
        "# FE Outer race 0.014 inches centered @6:00\n",
        "acquisitions[\"FEOR.014@6_0\"] = \"313.mat\"\n",
        "# FE Outer race 0.021 inches centered @6:00\n",
        "acquisitions[\"FEOR.021@6_0\"] = \"315.mat\"\n",
        "# FE Outer race 0.021 inches centered @3:00\n",
        "acquisitions[\"FEOR.021@3_1\"] = \"316.mat\"\n",
        "acquisitions[\"FEOR.021@3_2\"] = \"317.mat\"\n",
        "acquisitions[\"FEOR.021@3_3\"] = \"318.mat\"\n",
        "# DE Inner Race 0.028 inches\n",
        "acquisitions[\"DEIR.028_0\"] = \"3001.mat\"\n",
        "acquisitions[\"DEIR.028_1\"] = \"3002.mat\"\n",
        "acquisitions[\"DEIR.028_2\"] = \"3003.mat\"\n",
        "acquisitions[\"DEIR.028_3\"] = \"3004.mat\"\n",
        "# DE Ball 0.028 inches\n",
        "acquisitions[\"DEB.028_0\"] = \"3005.mat\"\n",
        "acquisitions[\"DEB.028_1\"] = \"3006.mat\"\n",
        "acquisitions[\"DEB.028_2\"] = \"3007.mat\"\n",
        "acquisitions[\"DEB.028_3\"] = \"3008.mat\""
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hQu5HBLioPIu",
        "colab_type": "text"
      },
      "source": [
        "# Functions definitions\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "YcKzJQ-XoSCg",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def get_labels_dict(acquisitions, separator='_', detectPosition=True):\n",
        "  \"\"\"Generate a dictionary linking the labels with values to keep consistence.\"\"\"\n",
        "  labels_dict = {}\n",
        "  value = 0\n",
        "  for key in acquisitions.keys():\n",
        "    key = key.split('_')[0]\n",
        "    key = key.split(separator)\n",
        "    if key[0] == \"Normal\" or detectPosition:\n",
        "      label = key[0]\n",
        "    else:\n",
        "      label = key[0][2:]\n",
        "    if not label in labels_dict:\n",
        "      labels_dict[label] = value\n",
        "      value += 1\n",
        "  return labels_dict"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bL85KQ_noXAV",
        "colab_type": "text"
      },
      "source": [
        "Convert Matlab file into tensors.\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "i52OZlbIodym",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "import scipy.io\n",
        "import numpy as np\n",
        "def acquisition2tensor(file_name, position=None, sample_size=sample_size):\n",
        "  \"\"\"\n",
        "  Convert Matlab file into tensors.\n",
        "  The file is divided in segments of sample_size values.\n",
        "  \"\"\"\n",
        "  print(file_name, end=' ')\n",
        "  matlab_file = scipy.io.loadmat(file_name)\n",
        "  DE_samples = []\n",
        "  FE_samples = []\n",
        "  \n",
        "  #signal segmentation\n",
        "  signal_begin = 0\n",
        "  if position == None:\n",
        "    DE_time = [key for key in matlab_file if key.endswith(\"DE_time\")][0] #Find the DRIVE END acquisition key name\n",
        "    FE_time = [key for key in matlab_file if key.endswith(\"FE_time\")][0] #Find the FAN END acquisition key name\n",
        "    acquisition_size = max(len(matlab_file[DE_time]),len(matlab_file[FE_time]))\n",
        "    while signal_begin + sample_size < acquisition_size:\n",
        "      DE_samples.append([item for sublist in matlab_file[DE_time][signal_begin:signal_begin+sample_size] for item in sublist])\n",
        "      FE_samples.append([item for sublist in matlab_file[FE_time][signal_begin:signal_begin+sample_size] for item in sublist])\n",
        "      signal_begin += sample_size\n",
        "    sample_tensor = np.stack([DE_samples,FE_samples],axis=2).astype('float32')\n",
        "  elif position == 'DE':\n",
        "    DE_time = [key for key in matlab_file if key.endswith(\"DE_time\")][0] #Find the DRIVE END acquisition key name\n",
        "    acquisition_size = len(matlab_file[DE_time])\n",
        "    while signal_begin + sample_size < acquisition_size:\n",
        "      DE_samples.append([item for sublist in matlab_file[DE_time][signal_begin:signal_begin+sample_size] for item in sublist])\n",
        "      signal_begin += sample_size\n",
        "    sample_tensor = np.stack([DE_samples],axis=2).astype('float32')\n",
        "  elif position == 'FE':\n",
        "    FE_time = [key for key in matlab_file if key.endswith(\"FE_time\")][0] #Find the FAN END acquisition key name\n",
        "    acquisition_size = len(matlab_file[FE_time])\n",
        "    while signal_begin + sample_size < acquisition_size:\n",
        "      FE_samples.append([item for sublist in matlab_file[FE_time][signal_begin:signal_begin+sample_size] for item in sublist])\n",
        "      signal_begin += sample_size\n",
        "    sample_tensor = np.stack([FE_samples],axis=2).astype('float32')\n",
        "  return sample_tensor"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "GlnwDCj9olhH",
        "colab_type": "text"
      },
      "source": [
        "Extract datasets from acquisitions.\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "-xJlTxjTomWi",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def concatenate_datasets(xd,yd,xo,yo):\n",
        "  \"\"\"\n",
        "  xd: destination patterns tensor\n",
        "  yd: destination labels tensor\n",
        "  xo: origin patterns tensor to be concateneted \n",
        "  yo: origin labels tensor to be concateneted \n",
        "  \"\"\"\n",
        "  if xd is None or yd is None:\n",
        "    xd = xo\n",
        "    yd = yo\n",
        "  else:\n",
        "    xd = np.concatenate((xd,xo))\n",
        "    yd = np.concatenate((yd,yo))\n",
        "  return xd,yd\n",
        "\n",
        "import urllib.request\n",
        "import os.path\n",
        "\n",
        "def acquisitions_from_substr(substr, acquisitions, labels_dict, position=None,\n",
        "                             url=\"http://csegroups.case.edu/sites/default/files/bearingdatacenter/files/Datafiles/\"):\n",
        "  \"\"\"\n",
        "  Extract samples from all files with some load.\n",
        "  \"\"\"\n",
        "  samples = None\n",
        "  labels = None\n",
        "  for key in acquisitions:\n",
        "    if str(substr) in key:\n",
        "      file_name = acquisitions[key]\n",
        "      if not os.path.exists(file_name):\n",
        "        urllib.request.urlretrieve(url+file_name, file_name)\n",
        "      if substr[:2] == key[:2] and position == None:\n",
        "        acquisition_samples = acquisition2tensor(file_name)\n",
        "      elif position =='DE':\n",
        "        acquisition_samples = acquisition2tensor(file_name, 'DE')\n",
        "      elif position =='FE':\n",
        "        acquisition_samples = acquisition2tensor(file_name, 'FE')\n",
        "      else:\n",
        "        acquisition_samples = acquisition2tensor(file_name, key[:2])\n",
        "      for label in labels_dict.keys():\n",
        "        if label in key:\n",
        "          break\n",
        "      acquisition_labels = np.ones(acquisition_samples.shape[0])*labels_dict[label]\n",
        "      samples,labels = concatenate_datasets(samples,labels,acquisition_samples,acquisition_labels)\n",
        "  print(substr)\n",
        "  return samples,labels"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "h1bmT_Proq-y",
        "colab_type": "text"
      },
      "source": [
        "# Downloading and Matlab files\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "LOwEbwZPowH2",
        "colab_type": "text"
      },
      "source": [
        "Extract samples.\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "iZG8Cuylorth",
        "colab_type": "code",
        "outputId": "e5e81450-0bb8-44c8-ce7f-ee3bff174b36",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 258
        }
      },
      "source": [
        "labels_dict = get_labels_dict(acquisitions, '.', False)\n",
        "print(labels_dict)\n",
        "def normal_indenpendent_position_acquisitions(load,acquisitions,labels_dict):\n",
        "  x,y = None,None\n",
        "  for position in ['DE','FE']:\n",
        "    xn,yn = acquisitions_from_substr('Normal_'+str(load),acquisitions,labels_dict,position)\n",
        "    x,y = concatenate_datasets(x,y,xn,yn)\n",
        "  return x,y\n",
        "\n",
        "xn_0,yn_0 = normal_indenpendent_position_acquisitions(0,acquisitions,labels_dict)\n",
        "xn_1,yn_1 = normal_indenpendent_position_acquisitions(1,acquisitions,labels_dict)\n",
        "xn_2,yn_2 = normal_indenpendent_position_acquisitions(2,acquisitions,labels_dict)\n",
        "xn_3,yn_3 = normal_indenpendent_position_acquisitions(3,acquisitions,labels_dict)\n",
        "\n",
        "x007,y007 = acquisitions_from_substr('007',acquisitions,labels_dict)\n",
        "x014,y014 = acquisitions_from_substr('014',acquisitions,labels_dict)\n",
        "x021,y021 = acquisitions_from_substr('021',acquisitions,labels_dict)\n",
        "x028,y028 = acquisitions_from_substr('028',acquisitions,labels_dict)\n",
        "\n",
        "severities = ['007','014','021','028']"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "{'Normal': 0, 'IR': 1, 'B': 2, 'OR': 3}\n",
            "97.mat Normal_0\n",
            "97.mat Normal_0\n",
            "98.mat Normal_1\n",
            "98.mat Normal_1\n",
            "99.mat Normal_2\n",
            "99.mat Normal_2\n",
            "100.mat Normal_3\n",
            "100.mat Normal_3\n",
            "105.mat 106.mat 107.mat 108.mat 118.mat 119.mat 120.mat 121.mat 130.mat 131.mat 132.mat 133.mat 144.mat 145.mat 146.mat 147.mat 156.mat 158.mat 159.mat 160.mat 282.mat 283.mat 284.mat 285.mat 278.mat 279.mat 280.mat 281.mat 294.mat 295.mat 296.mat 297.mat 298.mat 299.mat 300.mat 301.mat 302.mat 305.mat 306.mat 307.mat 007\n",
            "169.mat 170.mat 171.mat 172.mat 185.mat 186.mat 187.mat 188.mat 197.mat 198.mat 199.mat 200.mat 274.mat 275.mat 276.mat 277.mat 286.mat 287.mat 288.mat 289.mat 310.mat 309.mat 311.mat 312.mat 313.mat 014\n",
            "222.mat 223.mat 224.mat 225.mat 270.mat 271.mat 272.mat 273.mat 209.mat 210.mat 211.mat 212.mat 234.mat 235.mat 236.mat 237.mat 246.mat 247.mat 248.mat 249.mat 258.mat 259.mat 260.mat 261.mat 290.mat 291.mat 292.mat 293.mat 315.mat 316.mat 317.mat 318.mat 021\n",
            "3001.mat 3002.mat 3003.mat 3004.mat 3005.mat 3006.mat 3007.mat 3008.mat 028\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "v4iA0dt9pXyK",
        "colab_type": "text"
      },
      "source": [
        "Count number of samples.\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "CEHJBo9bpZGQ",
        "colab_type": "code",
        "outputId": "3135c587-a075-4ab3-9f76-98fd99b8706b",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 187
        }
      },
      "source": [
        "print(\"Label\", end='\\t')\n",
        "for s in severities:\n",
        "  print(s, end='\\t')\n",
        "print(\"total\")\n",
        "mat = np.zeros((4,4))\n",
        "i = 0\n",
        "for label,value in labels_dict.items():\n",
        "  print(label, end='\\t')\n",
        "  tsamples = 0\n",
        "  if label == 'Normal':\n",
        "    print(4*'\\t'+'...')\n",
        "    for load in range(4):\n",
        "      print(' '+str((load+len(severities)+1)%4)+(load)*'\\t', end='\\t')\n",
        "      mat[i][load] = list(eval('yn_'+str((load+len(severities)+1)%4))).count(value)\n",
        "      print(int(mat[i][load]))\n",
        "  else:\n",
        "    for j,severity in enumerate(severities):\n",
        "      tmp = eval('y'+str(severity))\n",
        "      if tmp is not None:\n",
        "        nsamples = list(tmp).count(value)\n",
        "        mat[i][j] = nsamples\n",
        "        print(nsamples, end='\\t')\n",
        "        tsamples += nsamples\n",
        "      else:\n",
        "        print('0', end='\\t')\n",
        "    print(tsamples)\n",
        "  i+=1\n",
        "total = np.sum(mat,axis=0)\n",
        "print(\"Total:\", end='\\t')\n",
        "for i in range(len(total)):\n",
        "  print(int(total[i]), end='\\t')\n",
        "print(int(np.sum(total)))"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Label\t007\t014\t021\t028\ttotal\n",
            "Normal\t\t\t\t\t...\n",
            " 1\t1890\n",
            " 2\t\t1890\n",
            " 3\t\t\t1896\n",
            " 0\t\t\t\t952\n",
            "IR\t1900\t1895\t1894\t946\t6635\n",
            "B\t1894\t1900\t1897\t944\t6635\n",
            "OR\t5694\t2131\t3798\t0\t11623\n",
            "Total:\t11378\t7816\t9485\t2842\t31521\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BQutiWuupevT",
        "colab_type": "text"
      },
      "source": [
        "#Gerando CSV Files\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "P2p5W3ZvU1MW",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "import os\n",
        "\n",
        "cond_dict = {v: k for k, v in labels_dict.items()}\n",
        "\n",
        "def write_csv(severity, samples, labels):\n",
        "  if not os.path.exists(\"cwru\"):\n",
        "    os.makedirs(\"cwru\")\n",
        "  sevdir = 'cwru/'+severity\n",
        "  if not os.path.exists(sevdir):\n",
        "    os.makedirs(sevdir)\n",
        "  for i,value in enumerate(labels):\n",
        "    condir = sevdir+'/'+cond_dict[value]\n",
        "    if not os.path.exists(condir):\n",
        "      os.makedirs(condir)\n",
        "    sample_name = condir+'/'+str(i).zfill(len(str(labels.shape[0])))+'.csv'\n",
        "    np.savetxt(sample_name,samples[i],delimiter=',')"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ZAT4AJoGZJ7k",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def write_dataset(severity, normal_load):\n",
        "  x,y = concatenate_datasets(eval('x'+severity),\n",
        "                             eval('y'+severity),\n",
        "                             eval('xn_'+str(normal_load)),\n",
        "                             eval('yn_'+str(normal_load)))\n",
        "  write_csv(severity,x,y)\n",
        "\n",
        "write_dataset('007',1)\n",
        "write_dataset('014',2)\n",
        "write_dataset('021',3)\n",
        "write_dataset('028',0)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "aNQ5Wfs4X4On",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "import shutil\n",
        "output_filename = 'cwru_segmented'\n",
        "ext = 'zip'\n",
        "shutil.make_archive(output_filename, ext, 'cwru')\n",
        "zipfile_name = output_filename+'.'+ext"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Iv3iuhNOTZ5E",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "from google.colab import files\n",
        "import time\n",
        "while not os.path.exists(zipfile_name):\n",
        "  time.sleep(1)\n",
        "files.download(zipfile_name)"
      ],
      "execution_count": 0,
      "outputs": []
    }
  ]
}