load model

README.md

@@ -1,7 +1,7 @@
 ---
 title: Hertz
-emoji: 🐢
-colorFrom: yellow
+emoji: 🐙
+colorFrom: green
 colorTo: purple
 sdk: gradio
 sdk_version: 4.28.2

app.py

@@ -0,0 +1,106 @@
+import gradio as gr
+import os
+import shutil
+from utils import segment, plotter, writer
+
+
+class GlobalState:
+    """
+    Class to store global variables
+    """
+    heart_area = [.54, .5, -.14, .16]
+    smooth_factor = 3
+    video_file_path = os.path.join(os.path.dirname(__file__), 'videos/')
+    input_path = None
+    line = None
+    result_file_path = os.path.join(os.path.dirname(__file__), 'result/result.mp4')
+    result_folder = os.path.join(os.path.dirname(__file__), 'result/')
+    yolo_path = os.path.join(os.path.dirname(__file__), 'pretrained_yolo.pt')
+
+
+def upload_video_file(fid):
+    """
+    uploads and save video to workdir
+    """
+    raw_path = os.path.join(GlobalState.video_file_path, os.path.basename(fid.name))
+    shutil.move(fid.name, raw_path)
+    GlobalState.input_path = raw_path
+    gr.Info("Video uploaded")
+    return gr.update('Run!')
+
+
+def processing(sl):
+    GlobalState.smooth_factor = int(sl)
+    graph, frames, message = segment(GlobalState.input_path, GlobalState.yolo_path, start=0, fstep=1,
+                                     crop=GlobalState.heart_area)
+    gr.Info(message)
+    
+    if message == 'Video processing succeeded':
+        GlobalState.line = graph
+        writer(GlobalState.result_file_path, frames)
+        gr.Info('Processed video saved!')
+        return gr.update(visible=True), gr.update(visible=True)
+    else:
+        return gr.update(visible=False), gr.update(visible=False)
+    
+    
+def plot_graph(sl):
+    sl = int(sl)
+    result, text = plotter(GlobalState.line, sl)
+    return result, gr.update(value=text)
+
+def show_video(btn):
+    return gr.update(label="Segmented Echo", value=GlobalState.result_file_path)
+
+def main():
+
+    shutil.rmtree(os.path.join(os.path.dirname(__file__), 'videos/'), ignore_errors=True)
+    shutil.rmtree(os.path.join(os.path.dirname(__file__), 'result/'), ignore_errors=True)
+    os.mkdir(os.path.join(os.path.dirname(__file__), 'videos/'))
+    os.mkdir(os.path.join(os.path.dirname(__file__), 'result/'))
+
+    with gr.Blocks() as demo:
+        with gr.Tab("Load"):
+            with gr.Row():
+                gr.Markdown(
+                    """
+                    # Load video file 🫀
+                    # Then press **Run!**
+                    # Have fun:)
+                    """)
+            with gr.Row():
+                with gr.Column():
+                    with gr.Row():
+                        video_upload = gr.File(label="Upload heart Echo", file_types=["video"], file_count="single")
+                    with gr.Row():
+                        process_button = gr.Button("Run!")
+                    with gr.Row():
+                        player = gr.Video(label="Segmented Echo", value=None, format='mp4')
+
+                with gr.Column():
+                    with gr.Row():
+                        smoother = gr.Slider(1, 50, 5, 1, label="Rolling Mean Window")
+                    with gr.Row():
+                        messenger = gr.Textbox(label='Ejection Fracture', value=None)
+                    with gr.Row():
+                        plot = gr.LinePlot(x="Frame", y="Left ventricle visible area, px*px",
+                                           overlay_point=False, 
+                                           tooltip=["Frame", "Left ventricle visible area, px*px"], 
+                                           width=500, height=300)
+                    with gr.Row():
+                        show_graph = gr.Button('Plot', visible=False)
+                    with gr.Row():
+                        show_button = gr.Button("Show result!")
+
+        video_upload.upload(upload_video_file, video_upload, outputs=[process_button], show_progress='full')
+        process_button.click(processing, inputs=[smoother], outputs=[show_graph, show_button], show_progress='full')
+        show_graph.click(plot_graph, inputs=[smoother], outputs=[plot, messenger])
+        show_button.click(show_video, outputs=[player])
+        player.change(show_video, outputs=[player])
+
+
+    demo.launch(allowed_paths=[GlobalState.result_folder])
+
+
+if __name__ == "__main__":
+    main()

pretrained_yolo.pt

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

requirements.txt

@@ -0,0 +1,8 @@
+opencv-python
+tqdm
+matplotlib
+ultralytics
+numpy
+pandas
+gradio
+ffmpeg-python

utils.py

@@ -0,0 +1,100 @@
+import ultralytics
+from ultralytics import YOLO
+import torch
+import numpy as np
+import pandas as pd
+from tqdm import tqdm
+from matplotlib import pyplot as plt
+import cv2
+import warnings
+import ffmpeg
+
+warnings.filterwarnings("ignore")
+device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
+
+
+
+
+def segment(video_path: str, model_path: str, start: int, fstep: int, crop: list) -> tuple:
+    """
+    runs YOLO segmentation model and calculates the LV Area
+    """
+    model = YOLO(model_path)#.to(device)
+    cap = cv2.VideoCapture(video_path)
+    stop = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
+    lv_area = []
+    frames = []
+    message = 'Video processing succeeded'
+    for fr in tqdm(range(start, stop, fstep), desc=f'processing ECHO'):
+        cap.set(cv2.CAP_PROP_POS_FRAMES, fr)
+        _, frame = cap.read()
+        new_w = int(frame.shape[1] * crop[1])
+        new_h = int(frame.shape[0] * crop[0])
+        new_left = int(frame.shape[1] / 2 + crop[3] * new_w - new_w / 2)
+        new_top = int(frame.shape[0] / 2 + crop[2] * new_h - new_h / 2)
+        frame = frame[new_top:new_top + new_h, new_left:new_left + new_w]
+        frame_m = frame
+        inputs = frame #torch.Tensor(frame).to(device)
+        with torch.no_grad():
+            result = model(inputs, verbose=False)
+        result = result#.to('cpu')
+        classes = result[0].names
+        if len(classes) == 0:
+            pass
+        overlay = frame.copy()
+        color_list = [(255, 0, 0),
+                      (255, 255, 0),
+                      (255, 0, 255),
+                      (0, 255, 0),
+                      (0, 0, 255),
+                      (128, 128, 128)]
+        for i, res in enumerate(result[0]):
+            bx = res.boxes
+            m = res.masks.xy
+            label = int(bx.cls.squeeze().cpu())
+            if label == 1:
+                lv_area.append(cv2.contourArea(m[0]))
+            box = list(map(int, bx.xyxy.squeeze().cpu().tolist()))
+            cv2.rectangle(overlay, (box[0], box[1]), (box[2], box[3]), (36, 255, 12), 2)
+            cv2.putText(overlay, classes[label], (box[0], box[1] - 5), cv2.FONT_HERSHEY_TRIPLEX, 1, (0, 0, 255), 2)
+            cv2.fillPoly(overlay, pts=np.int32([m]), color=color_list[i % 6])
+            alpha = 0.4
+            frame_m = cv2.addWeighted(overlay, alpha, frame, 1 - alpha, 0)
+        frames.append(frame_m)
+        if len(lv_area) == 0:
+            message = 'Video processing failed'
+    return lv_area, frames, message
+
+
+def plotter(lv_data: list, window: int) -> tuple:
+    """
+    plots the rolling mean graph for LV area.
+    calculates the average ejection fracture
+    """
+    lv_rolling = pd.Series(lv_data).rolling(window=window).mean().dropna()
+    ef = (max(lv_rolling) - min(lv_rolling)) / max(lv_rolling)
+    dataframe = pd.DataFrame({
+      'Frame': np.array(range(len(lv_rolling))),
+      'Left ventricle visible area, px*px': lv_rolling.values
+    }).astype('int32')
+    txt = f'Ejection fraction - {ef:.1%}'
+    return dataframe, txt
+    
+
+def writer(fn, images, framerate=25, vcodec='libx264'):
+    if not isinstance(images, np.ndarray):
+        images = np.asarray(images)
+    n, height, width, channels = images.shape
+    process = (
+        ffmpeg
+            .input('pipe:', format='rawvideo', pix_fmt='rgb24', s='{}x{}'.format(width, height))
+            .output(fn, pix_fmt='yuv420p', vcodec=vcodec, r=framerate)
+            .overwrite_output()
+            .run_async(pipe_stdin=True)
+    )
+    for frame in images:
+        process.stdin.write(
+            frame.astype(np.uint8).tobytes()
+        )
+    process.stdin.close()
+    process.wait()