Upload 8 files

README.md

@@ -1,12 +1,24 @@
----
-title: HearNet
-emoji: 🔥
-colorFrom: purple
-colorTo: gray
-sdk: streamlit
-sdk_version: 1.17.0
-app_file: app.py
-pinned: false
----
-
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# HeartNet
+
+A joint project by [oapostrophe](https://github.com/oapostrophe), [gkenderova](https://github.com/gkenderova), [soksamnanglim](https://github.com/soksamnanglim), [syaa2018](https://github.com/syaa2018)
+
+For a high-level overview of this project, check out this [blog post](https://oapostrophe.github.io/heartnet/) and [90-second demo](https://www.youtube.com/watch?v=EqAU-FRu6C4).  For a full presentation and more detailed writeup on our methodology, check out the report on our [project website](https://oapostrophe.github.io/HeartNet/).
+
+The trained model can be demoed by downloading `app.py` and `demo_model.pkl`, installing [streamlit](https://anaconda.org/conda-forge/streamlit) and [fastai](https://pypi.org/project/fastai/), then running:
+```shell
+streamlit run app.py
+```
+You can then visit the provided url in your browser; for convenience, sample generated MI and Normal EKG images are provided in the `/test files` directory.
+
+To use any of the other files, you'll have to download the [PTB-XL](https://physionet.org/content/ptb-xl/1.0.1/) dataset.
+
+The important files are the following:
+- `app.py` StreamLit-based web interface using a trained model
+- `dataset generation/generate_imgset1.py` our first iteration generating a dataset directly with MatPlotLib; these images look rough.
+- `dataset generation/generate_imgset2.py` our second iteration that generates nicer-looking images
+- `dataset generation/generate_imgset3.py` adds random simulated shadows overlaying generated images
+- `dataset generation/generate_rnn_imgset.py` generates individual images for each of 12 leads, for input into an RNN (rnn code currently fails to learn).
+- `dataset generation/automold.py` library with image augmentation code for adding shadows
+- `training/cnn_learner.py` trains and saves a cnn on generated images.
+
+Feel free to [email me](swow2015@mymail.pomona.edu) with any questions!

app.py

@@ -0,0 +1,45 @@
+from fastai.vision.all import *
+from io import BytesIO
+import requests
+import streamlit as st
+
+"""
+# HeartNet
+This is a classifier for images of 12-lead EKGs.  It will attempt to detect whether the EKG indicates an acute MI.  It was trained on simulated images.
+"""
+
+def predict(img):
+    st.image(img, caption="Your image", use_column_width=True)
+    pred, _, probs = learn_inf.predict(img)
+    # st.write(learn_inf.predict(img))
+
+    f"""
+    ## This **{'is ' if pred == 'mi' else 'is not'}** an MI (heart attack).
+    ### Probability of MI: {probs[0].item()*100: .2f}%
+    ### Probability Normal: {probs[1].item()*100: .2f}%
+    """
+
+
+path = "./"
+learn_inf = load_learner(path + "demo_model.pkl")
+
+option = st.radio("", ["Upload Image", "Image URL"])
+
+if option == "Upload Image":
+    uploaded_file = st.file_uploader("Please upload an image.")
+
+    if uploaded_file is not None:
+        img = PILImage.create(uploaded_file)
+        predict(img)
+
+else:
+    url = st.text_input("Please input a url.")
+
+    if url != "":
+        try:
+            response = requests.get(url)
+            pil_img = PILImage.create(BytesIO(response.content))
+            predict(pil_img)
+
+        except:
+            st.text("Problem reading image from", url)

demo_model.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a787a70f3c2197c26beafd169d2e6b3364414fec7066eb7468a794d0a209ea68
+size 50315747

training/cnn_learner.py

@@ -0,0 +1,31 @@
+import fastbook
+fastbook.setup_book()
+from fastbook import *
+from fastai.vision.widgets import *
+
+# Pick a GPU with free resources (change this accordingly)
+torch.cuda.set_device(0)
+
+# Get images
+image_path = Path('/raid/heartnet/data/imgset2')
+images = get_image_files(image_path)
+
+# Initialize metric functions
+recall_function = Recall(pos_label=0)
+precision_function = Precision(pos_label=0)
+f1_score = F1Score(pos_label=0)
+
+# Initialize DataLoader
+images_datablock = DataBlock(
+    blocks=(ImageBlock, CategoryBlock), 
+    get_items=get_image_files, 
+    splitter=RandomSplitter(valid_pct=0.2, seed=42),
+    get_y=parent_label,
+    batch_tfms=aug_transforms(do_flip=False)
+)
+dls = images_datablock.dataloaders(image_path, bs=16)
+
+# Create, train, and save model
+learn = cnn_learner(dls, resnet152, metrics=[error_rate, recall_function, precision_function, f1_score])
+learn.fine_tune(16)
+learn.export('demo_model_50.pkl')

training/rnn (does not work)/ConvLSTM.py

@@ -0,0 +1,194 @@
+import torch.nn as nn
+import torch
+from torchvision import transforms
+from PIL import Image
+from pathlib import Path
+import os
+import shutil
+from fastai.vision.all import *
+
+class ConvLSTMCell(nn.Module):
+
+    def __init__(self, input_dim, hidden_dim, kernel_size, bias):
+        """
+        Initialize ConvLSTM cell.
+        Parameters
+        ----------
+        input_dim: int
+            Number of channels of input tensor.
+        hidden_dim: int
+            Number of channels of hidden state.
+        kernel_size: (int, int)
+            Size of the convolutional kernel.
+        bias: bool
+            Whether or not to add the bias.
+        """
+
+        super(ConvLSTMCell, self).__init__()
+
+        self.input_dim = input_dim
+        self.hidden_dim = hidden_dim
+
+        self.kernel_size = kernel_size
+        self.padding = kernel_size[0] // 2, kernel_size[1] // 2
+        self.bias = bias
+
+        self.conv = nn.Conv2d(in_channels=self.input_dim + self.hidden_dim,
+                              out_channels=4 * self.hidden_dim,
+                              kernel_size=self.kernel_size,
+                              padding=self.padding,
+                              bias=self.bias)
+
+    def forward(self, input_tensor, cur_state):
+        h_cur, c_cur = cur_state
+
+        combined = torch.cat([input_tensor, h_cur], dim=1)  # concatenate along channel axis
+
+        combined_conv = self.conv(combined)
+        cc_i, cc_f, cc_o, cc_g = torch.split(combined_conv, self.hidden_dim, dim=1)
+        i = torch.sigmoid(cc_i)
+        f = torch.sigmoid(cc_f)
+        o = torch.sigmoid(cc_o)
+        g = torch.tanh(cc_g)
+
+        c_next = f * c_cur + i * g
+        h_next = o * torch.tanh(c_next)
+
+        return h_next, c_next
+
+    def init_hidden(self, batch_size, image_size):
+        height, width = image_size
+        return (torch.zeros(batch_size, self.hidden_dim, height, width, device=self.conv.weight.device),
+                torch.zeros(batch_size, self.hidden_dim, height, width, device=self.conv.weight.device))
+
+
+class ConvLSTM(nn.Module):
+
+    """
+    Parameters:
+        input_dim: Number of channels in input
+        hidden_dim: Number of hidden channels
+        kernel_size: Size of kernel in convolutions
+        num_layers: Number of LSTM layers stacked on each other
+        batch_first: Whether or not dimension 0 is the batch or not
+        bias: Bias or no bias in Convolution
+        return_all_layers: Return the list of computations for all layers
+        Note: Will do same padding.
+    Input:
+        A tensor of size B, T, C, H, W or T, B, C, H, W
+    Output:
+        A tuple of two lists of length num_layers (or length 1 if return_all_layers is False).
+            0 - layer_output_list is the list of lists of length T of each output
+            1 - last_state_list is the list of last states
+                    each element of the list is a tuple (h, c) for hidden state and memory
+    Example:
+        >> x = torch.rand((32, 10, 64, 128, 128))
+        >> convlstm = ConvLSTM(64, 16, 3, 1, True, True, False)
+        >> _, last_states = convlstm(x)
+        >> h = last_states[0][0]  # 0 for layer index, 0 for h index
+    """
+
+    def __init__(self, input_dim, hidden_dim, kernel_size, num_layers,
+                 batch_first=False, bias=True, return_all_layers=False):
+        super(ConvLSTM, self).__init__()
+
+        self._check_kernel_size_consistency(kernel_size)
+
+        # Make sure that both `kernel_size` and `hidden_dim` are lists having len == num_layers
+        kernel_size = self._extend_for_multilayer(kernel_size, num_layers)
+        hidden_dim = self._extend_for_multilayer(hidden_dim, num_layers)
+        if not len(kernel_size) == len(hidden_dim) == num_layers:
+            raise ValueError('Inconsistent list length.')
+
+        self.input_dim = input_dim
+        self.hidden_dim = hidden_dim
+        self.kernel_size = kernel_size
+        self.num_layers = num_layers
+        self.batch_first = batch_first
+        self.bias = bias
+        self.return_all_layers = return_all_layers
+
+        cell_list = []
+        for i in range(0, self.num_layers):
+            cur_input_dim = self.input_dim if i == 0 else self.hidden_dim[i - 1]
+
+            cell_list.append(ConvLSTMCell(input_dim=cur_input_dim,
+                                          hidden_dim=self.hidden_dim[i],
+                                          kernel_size=self.kernel_size[i],
+                                          bias=self.bias))
+
+        self.cell_list = nn.ModuleList(cell_list)
+
+    def forward(self, input_tensor, hidden_state=None):
+        """
+        Parameters
+        ----------
+        input_tensor: todo
+            5-D Tensor either of shape (t, b, c, h, w) or (b, t, c, h, w)
+        hidden_state: todo
+            None. todo implement stateful
+        Returns
+        -------
+        last_state_list, layer_output
+        """
+        if not self.batch_first:
+            # (t, b, c, h, w) -> (b, t, c, h, w)
+            input_tensor = input_tensor.permute(1, 0, 2, 3, 4)
+
+        b, _, _, h, w = input_tensor.size()
+
+        # Implement stateful ConvLSTM
+        if hidden_state is not None:
+            raise NotImplementedError()
+        else:
+            # Since the init is done in forward. Can send image size here
+            hidden_state = self._init_hidden(batch_size=b,
+                                             image_size=(h, w))
+
+        layer_output_list = []
+        last_state_list = []
+
+        seq_len = input_tensor.size(1)
+        cur_layer_input = input_tensor
+
+        for layer_idx in range(self.num_layers):
+
+            h, c = hidden_state[layer_idx]
+            output_inner = []
+            for t in range(seq_len):
+                h, c = self.cell_list[layer_idx](input_tensor=cur_layer_input[:, t, :, :, :],
+                                                 cur_state=[h, c])
+                output_inner.append(h)
+
+            layer_output = torch.stack(output_inner, dim=1)
+            cur_layer_input = layer_output
+
+            layer_output_list.append(layer_output)
+            last_state_list.append([h, c])
+
+        if not self.return_all_layers:
+            layer_output_list = layer_output_list[-1:]
+            last_state_list = last_state_list[-1:]
+
+        return layer_output_list, last_state_list
+
+    def _init_hidden(self, batch_size, image_size):
+        init_states = []
+        for i in range(self.num_layers):
+            init_states.append(self.cell_list[i].init_hidden(batch_size, image_size))
+        return init_states
+
+    @staticmethod
+    def _check_kernel_size_consistency(kernel_size):
+        if not (isinstance(kernel_size, tuple) or
+                (isinstance(kernel_size, list) and all([isinstance(elem, tuple) for elem in kernel_size]))):
+            raise ValueError('`kernel_size` must be tuple or list of tuples')
+
+    @staticmethod
+    def _extend_for_multilayer(param, num_layers):
+        if not isinstance(param, list):
+            param = [param] * num_layers
+        return param
+
+
+

training/rnn (does not work)/generate_mi_tensors.py

@@ -0,0 +1,37 @@
+import torch.nn as nn
+import torch
+from torchvision import transforms
+from PIL import Image
+from pathlib import Path
+import os
+import shutil
+
+
+mi_src = "./imgset_rnn/mi/"
+mi_dest = "./tensorfiles_rnn/mi"
+
+tensornum = 0
+
+for dirs, sub_dirs, files in os.walk(mi_src):
+    if (dirs != "./imgset_rnn/mi/"):
+        tensors = []
+        pathlist = Path(dirs).glob('**/*.png')
+        for path in pathlist:
+            path_in_str = str(path) # because path is object not string
+            #print(path_in_str)
+            
+            pil_img = Image.open(path_in_str).convert("RGB")
+            pil_to_tensor = transforms.ToTensor()(pil_img)
+            tensors.append(pil_to_tensor)
+
+        tensor = torch.stack(tuple(tensors))
+        tensor2 = torch.unsqueeze(tensor, 0)
+        
+        tensorfile = torch.save(tensor2, 'tensor' + str(tensornum) + '.pt')
+  
+        shutil.move('tensor' + str(tensornum) + '.pt', './tensorfiles_rnn/mi/tensor' + str(tensornum) + '.pt') 
+        print('moved ' + str(tensornum) + '!')
+        tensornum += 1
+
+
+#generated 5486 mi tensors :) 

training/rnn (does not work)/generate_norm_tensors.py

@@ -0,0 +1,37 @@
+import torch.nn as nn
+import torch
+from torchvision import transforms
+from PIL import Image
+from pathlib import Path
+import os
+import shutil
+
+
+norm_src = "./imgset_rnn/normal/"
+norm_dest = "./tensorfiles_rnn/norm"
+
+tensornum = 0
+
+for dirs, sub_dirs, files in os.walk(norm_src):
+    if (dirs != "./imgset_rnn/normal/"):
+        tensors = []
+        pathlist = Path(dirs).glob('**/*.png')
+        for path in pathlist:
+            path_in_str = str(path) # because path is object not string
+            #print(path_in_str)
+            
+            pil_img = Image.open(path_in_str).convert("RGB")
+            pil_to_tensor = transforms.ToTensor()(pil_img)
+            tensors.append(pil_to_tensor)
+
+        tensor = torch.stack(tuple(tensors))
+        tensor2 = torch.unsqueeze(tensor, 0)
+        
+        tensorfile = torch.save(tensor2, 'tensor' + str(tensornum) + '.pt')
+  
+        shutil.move('tensor' + str(tensornum) + '.pt', './tensorfiles_rnn/norm/tensor' + str(tensornum) + '.pt') 
+        print('moved ' + str(tensornum) + '!')
+        tensornum += 1
+
+
+#generated 7547 tensors :) 

training/rnn (does not work)/rnn.py

@@ -0,0 +1,134 @@
+import torch
+from torch.utils.data import Dataset
+from torchvision import datasets
+from torchvision.transforms import ToTensor, Lambda
+import matplotlib.pyplot as plt
+import os
+import pandas as pd
+from torchvision.io import read_image
+from torch.utils.data import DataLoader
+from pathlib import Path
+import ConvLSTM
+import torch.optim as optim
+
+class CustomImageDataset(Dataset):
+    def __init__(self, image_directory): 
+        self.image_directory = image_directory
+
+        self.total_num_inputs = 0
+
+        for dir_name in Path(image_directory).glob('*'):
+            self.total_num_inputs += len(list(dir_name.glob('*.pt')))
+
+        self.tensor_labels = ["norm", "mi"]
+
+    def __len__(self):
+        return self.total_num_inputs
+
+    def __getitem__(self, idx):
+        #there are 7547 norm images and 5486 mi images
+        #return the correct label and the corresponding tensor (loaded by the file!)
+
+        if idx <= 7546:
+            label = self.tensor_labels[0]
+            tensor_file_path = "./tensorfiles_rnn/norm/tensor" + str(idx) + ".pt"
+            tensor = torch.load(tensor_file_path) 
+        else:
+            label = self.tensor_labels[1]
+            tensor_file_path = "./tensorfiles_rnn/mi/tensor" + str(idx-7547) + ".pt"
+            tensor = torch.load(tensor_file_path)
+
+        return (tensor.squeeze(), torch.tensor(0.0 if label == self.tensor_labels[0] else 1.0))
+
+training_data = CustomImageDataset("./tensorfiles_rnn")
+
+train_dataloader = DataLoader(training_data, batch_size=2, shuffle=True)
+
+# Display image and label.
+train_features, train_labels = next(iter(train_dataloader))
+
+print(f"Feature batch shape: {train_features.size()}")
+print(f"Labels batch shape: {train_labels.size()}")
+
+plt.imshow(train_features[0].squeeze()[0].squeeze().permute(1, 2, 0))
+
+
+
+class FlatConvLSTM(torch.nn.Module):
+    """An ConvLSTM layer that ignores the current hidden and cell states."""
+    def __init__(self):
+        super().__init__()
+        self.convlstm = ConvLSTM.ConvLSTM(3, 10, (3,3), 1, True, True, False)
+
+    def forward(self, x):
+        _, lstm_output = self.convlstm(x)
+        return lstm_output[0][0]
+
+
+model = torch.nn.Sequential(
+    FlatConvLSTM(),
+    torch.nn.Flatten(),
+    torch.nn.Linear(10*480*640, 1),
+    torch.nn.Sigmoid()
+)
+
+
+#create the loss function
+loss = torch.nn.BCELoss()
+optimizer = optim.Adam(model.parameters())
+device = torch.device("cuda:2" if torch.cuda.is_available() else "cpu")
+
+num_epochs = 10
+model.to(device)
+
+for epoch in range(num_epochs):
+    
+    # Set model to training mode
+    model.train()
+    
+    # Update the model for each batch
+    train_count = 0
+    train_cost = 0
+    batch = 0
+
+    for X, y in train_dataloader:
+        
+        # Compute model cost
+        #yhat = model(X.view(-1, nx))
+        X = X.to(device)
+        y = y.to(device)
+
+        yhat = model(X)
+        #print(yhat.shape, y.shape)
+
+        try:
+            cost = loss(yhat.squeeze(), y)
+            model.zero_grad()
+            cost.backward()
+        except:
+            print()
+
+        
+        train_count += X.shape[0]
+        train_cost += cost.item()
+        optimizer.step()
+        print(epoch, batch, cost.item())
+        batch += 1
+
+    # Set model to evaluation mode
+    model.eval()
+    
+    # Test model on validation data
+    valid_count = 0
+    valid_cost = 0
+    valid_correct = 0
+            
+    train_cost /= train_count
+    
+    print(epoch, train_cost)
+
+print('Done.')
+
+#save the model in a file
+torch.save(model, "./rnn_saved_models/model.py")
+torch.save(model.state_dict(), "./rnn_saved_models/model_parameters.py")