ecg plot

app.py

@@ -1,17 +1,34 @@
 import gradio as gr
 #from transformers import pipeline
 from transformers import AutoModel
-
+import ecg_plot
+import matplotlib.pyplot as plt
+from PIL import Image
 #pipeline = pipeline(task="image-classification", model="julien-c/hotdog-not-hotdog")
 model = AutoModel.from_pretrained("deepsynthbody/deepfake_ecg", trust_remote_code=True)
 
 def predict(num_ecgs):
-    predictions = model(int(num_ecgs))
-    return {"ecgs": predictions}
+    prediction = (model(1)[0].t()/1000) # to micro volte
+    
+    
+    lead_III = (prediction[1] - prediction[0]).unsqueeze(dim=0)
+    lead_aVR = ((prediction[0] + prediction[1])*(-0.5)).unsqueeze(dim=0)
+    lead_aVL = (prediction[0] -  prediction[1]* 0.5).unsqueeze(dim=0)
+    lead_aVF = (prediction[1] -  prediction[0]* 0.5).unsqueeze(dim=0)
+    all = torch.cat((prediction, lead_III, lead_aVR, lead_aVL, lead_aVF), dim=0)
+    all_corrected = all[torch.tensor([0,1,8, 9, 10, 11, 2,3,4,5,6,7])]
+    
+    ecg_plot.plot(all_corrected, sample_rate = 500, title = 'ECG 12')
+    
+    #ecg_plot.show()
+    buf = io.BytesIO()
+    plt.savefig(buf, format="png")
+    img = Image.open(buf)
+    return img
 
 gr.Interface(
     predict,
-    inputs="text",
-    outputs="text",
+    inputs=None,
+    outputs="image",
     title="Generating ECGs",
-).launch()
+).launch(share=True) 

requirements.txt

@@ -1,2 +1,5 @@
 transformers
-torch
+torch
+ecg-plot
+matplotlib
+PIL