Upload utils.py

utils.py

@@ -15,20 +15,21 @@ from scipy.signal import decimate, resample_poly, firwin, lfilter
 
 
 os.environ["CUDA_VISIBLE_DEVICES"]="0"
+device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
 
-def resample(signal, fs):
-    # downsample the signal to a sample rate of 256 Hz
-    if fs>256:
-        fs_down = 256 # Desired sample rate
+def resample(signal, fs, tgt_fs):
+    # downsample the signal to the target sample rate
+    if fs>tgt_fs:
+        fs_down = tgt_fs # Desired sample rate
         q = int(fs / fs_down) # Downsampling factor
         signal_new = []
         for ch in signal:
             x_down = decimate(ch, q)
             signal_new.append(x_down)
 
-    # upsample the signal to a sample rate of 256 Hz
-    elif fs<256:
-        fs_up = 256  # Desired sample rate
+    # upsample the signal to the target sample rate
+    elif fs<tgt_fs:
+        fs_up = tgt_fs  # Desired sample rate
         p = int(fs_up / fs)  # Upsampling factor 
         signal_new = []
         for ch in signal:
@@ -70,14 +71,14 @@ def cut_data(filepath, raw_data):
     total = int(len(raw_data[0]) / 1024)
     for i in range(total):
         table = raw_data[:, i * 1024:(i + 1) * 1024]
-        filename = filepath + '/temp2/' + str(i) + '.csv'
+        filename = filepath + 'temp2/' + str(i) + '.csv'
         with open(filename, 'w', newline='') as csvfile:
             writer = csv.writer(csvfile)
             writer.writerows(table)
     return total
 
 
-def glue_data(file_name, total, output):
+def glue_data(file_name, total):
     gluedata = 0
     for i in range(total):
         file_name1 = file_name + 'output{}.csv'.format(str(i))
@@ -96,11 +97,7 @@ def glue_data(file_name, total, output):
             raw_data[:, 1] = smooth
             gluedata = np.append(gluedata, raw_data, axis=1)
     #print(gluedata.shape)
-    filename2 = output
-    with open(filename2, 'w', newline='') as csvfile:
-        writer = csv.writer(csvfile)
-        writer.writerows(gluedata)
-        #print("GLUE DONE!" + filename2)
+    return gluedata
 
 
 def save_data(data, filename):
@@ -112,91 +109,105 @@ def dataDelete(path):
     try:
         shutil.rmtree(path)
     except OSError as e:
-        print(e)
+        pass
+        #print(e)
     else:
         pass
         #print("The directory is deleted successfully")
 
 
 def decode_data(data, std_num, mode=5):
-    decode = data
+    
     if mode == "ICUNet":
         # 1. read name
-        model = cumbersome_model2.UNet1(n_channels=30, n_classes=30)
+        model = cumbersome_model2.UNet1(n_channels=30, n_classes=30).to(device)
         resumeLoc = './model/ICUNet/modelsave' + '/checkpoint.pth.tar'
         # 2. load model
-        checkpoint = torch.load(resumeLoc, map_location='cpu')
+        checkpoint = torch.load(resumeLoc, map_location=device)
         model.load_state_dict(checkpoint['state_dict'], False)
         model.eval()
         # 3. decode strategy
         with torch.no_grad():
             data = data[np.newaxis, :, :]
-            data = torch.Tensor(data)
+            data = torch.Tensor(data).to(device)
             decode = model(data)
 
       
-    elif mode == "UNetpp" or mode == "AttUnet":
+    elif mode == "ICUNet++" or mode == "ICUNet_attn":
         # 1. read name
-        if mode == "UNetpp":
-            model = UNet_family.NestedUNet3(num_classes=30)
-        elif mode == "AttUnet":
-            model = UNet_attention.UNetpp3_Transformer(num_classes=30)
-        resumeLoc = './model/'+ mode + '/modelsave' + '/checkpoint.pth.tar'
+        if mode == "ICUNet++":
+            model = UNet_family.NestedUNet3(num_classes=30).to(device)
+        elif mode == "ICUNet_attn":
+            model = UNet_attention.UNetpp3_Transformer(num_classes=30).to(device)
+        resumeLoc = './model/' + mode + '/modelsave' + '/checkpoint.pth.tar'
         # 2. load model
-        checkpoint = torch.load(resumeLoc, map_location='cpu')
+        checkpoint = torch.load(resumeLoc, map_location=device)
         model.load_state_dict(checkpoint['state_dict'], False)
         model.eval()
         # 3. decode strategy
         with torch.no_grad():
             data = data[np.newaxis, :, :]
-            data = torch.Tensor(data)
+            data = torch.Tensor(data).to(device)
             decode1, decode2, decode = model(data)
             
 
-    elif mode == "EEGART":
+    elif mode == "ART":
         # 1. read name
         resumeLoc = './model/' + mode + '/modelsave/checkpoint.pth.tar'
         # 2. load model
-        checkpoint = torch.load(resumeLoc, map_location='cpu')
-        model = tf_model.make_model(30, 30, N=2)
+        checkpoint = torch.load(resumeLoc, map_location=device)
+        model = tf_model.make_model(30, 30, N=2).to(device)
         model.load_state_dict(checkpoint['state_dict'])
         model.eval()
         # 3. decode strategy
         with torch.no_grad():
-            data = torch.FloatTensor(data)
+            data = torch.FloatTensor(data).to(device)
             data = data.unsqueeze(0)
             src = data
-            tgt = data
+            tgt = data # you can modify to randomize data
             batch = tf_data.Batch(src, tgt, 0)
             out = model.forward(batch.src, batch.src[:,:,1:], batch.src_mask, batch.trg_mask)
             decode = model.generator(out)
             decode = decode.permute(0, 2, 1)
-            #add_tensor = torch.zeros(1, 30, 1)
-            #decode = torch.cat((decode, add_tensor), dim=2)
+            add_tensor = torch.zeros(1, 30, 1).to(device)
+            decode = torch.cat((decode, add_tensor), dim=2)
 
     # 4. numpy
     #print(decode.shape)
-    #decode = np.array(decode.cpu()).astype(np.float64)
-    decode = np.array(decode).astype(np.float64)
-    print(type(decode))
-    print(decode.shape)
-    #decode = decode.tolist()
+    decode = np.array(decode.cpu()).astype(np.float64)
     return decode
 
-def preprocessing(filepath, filename, samplerate):
+
+def reorder_data(raw_data, mapping_result):
+    new_data = np.zeros((30, raw_data.shape[1]))
+    zero_arr = np.zeros((1, raw_data.shape[1]))
+    for i, (indices, flag) in enumerate(zip(mapping_result["index"], mapping_result["isOriginalData"])):
+        if flag == True:
+            new_data[i, :] = raw_data[indices[0], :]
+        elif indices[0] == None:
+            new_data[i, :] = zero_arr
+        else:
+            data = [raw_data[idx, :] for idx in indices]
+            new_data[i, :] = np.mean(data, axis=0)
+    return new_data
+
+def preprocessing(filepath, inputfile, samplerate, mapping_result):
     # establish temp folder
     try:
-        os.mkdir(filepath+"/temp2/")
+        os.mkdir(filepath+"temp2/")
     except OSError as e:
-        dataDelete(filepath+"/temp2/")
-        os.mkdir(filepath+"/temp2/")
+        dataDelete(filepath+"temp2/")
+        os.mkdir(filepath+"temp2/")
         print(e)
     
     # read data
-    signal = read_train_data(filepath+'/'+filename)
+    signal = read_train_data(inputfile)
+    #print(signal.shape)
+    # channel mapping
+    signal = reorder_data(signal, mapping_result)
     #print(signal.shape)
     # resample
-    signal = resample(signal, samplerate)
+    signal = resample(signal, samplerate, 256)
     #print(signal.shape)
     # FIR_filter
     signal = FIR_filter(signal, 1, 50)
@@ -206,13 +217,29 @@ def preprocessing(filepath, filename, samplerate):
 
     return total_file_num
 
+def restore_order(data, all_data, mapping_result):
+    for i, (indices, flag) in enumerate(zip(mapping_result["index"], mapping_result["isOriginalData"])):
+        if flag == True:
+            all_data[indices[0], :] = data[i, :]
+    return all_data
+
+def postprocessing(data, samplerate, outputfile, mapping_result, batch_cnt, channel_num):
+    
+    # resample to original sampling rate
+    data = resample(data, 256, samplerate)
+    # reverse channel mapping
+    all_data = np.zeros((channel_num, data.shape[1])) if batch_cnt==0 else read_train_data(outputfile)
+    all_data = restore_order(data, all_data, mapping_result)
+    # save data
+    save_data(all_data, outputfile)
+
 
 # model = tf.keras.models.load_model('./denoise_model/')
-def reconstruct(model_name, total, filepath, outputfile):
+def reconstruct(model_name, total, filepath, batch_cnt):
     # -------------------decode_data---------------------------
     second1 = time.time()
     for i in range(total):
-        file_name = filepath + '/temp2/{}.csv'.format(str(i))
+        file_name = filepath + 'temp2/{}.csv'.format(str(i))
         data_noise = read_train_data(file_name)
         
         std = np.std(data_noise)
@@ -222,18 +249,17 @@ def reconstruct(model_name, total, filepath, outputfile):
 
         # Deep Learning Artifact Removal
         d_data = decode_data(data_noise, std, model_name)
-        #d_data = d_data[0]
+        d_data = d_data[0]
 
-        outputname = filepath + '/temp2/output{}.csv'.format(str(i))
+        outputname = filepath + 'temp2/output{}.csv'.format(str(i))
         save_data(d_data, outputname)
-        #d_data.to_csv(outputname, index=False)
-
 
     # --------------------glue_data----------------------------
-    glue_data(filepath+"/temp2/", total, filepath+'/'+outputfile)
+    data = glue_data(filepath+"temp2/", total)
     # -------------------delete_data---------------------------
-    dataDelete(filepath+"/temp2/")
+    dataDelete(filepath+"temp2/")
     second2 = time.time()
-
-    print("Using", model_name,"model to reconstruct", outputfile, " has been success in", second2 - second1, "sec(s)")
-
+    
+    print(f"Using {model_name} model to reconstruct batch-{batch_cnt+1} has been success in {second2 - second1} sec(s)")
+    return data
+