from pathlib import Path import sys import time from time import perf_counter import argparse from loguru import logger import os from predict import Model from datetime import datetime from scipy import signal import plotly.graph_objects as go import numpy as np import io from PIL import Image import cv2 from PySide6.QtCore import Qt, QThread, Signal, Slot from PySide6.QtGui import QImage, QPixmap from PySide6.QtWidgets import ( QApplication, QHBoxLayout, QLabel, QMainWindow, QPushButton, QSizePolicy, QVBoxLayout, QWidget, ) # for telemed import matplotlib.pyplot as plt import ctypes from ctypes import * # 720p video_w = 1280 video_h = 720 # Copy from detection.py from telemed sample code class Telemed: def __init__(self): # starting copy from the origianl main # Setting ultrasound size # w = 512 # h = 512 w = 640 h = 640 # Load dll # usgfw2 = cdll.LoadLibrary('./usgfw2wrapper_C++_sources/usgfw2wrapper/x64/Release/usgfw2wrapper.dll') usgfw2 = cdll.LoadLibrary("./usgfw2wrapper.dll") # Ultrasound initialize usgfw2.on_init() ERR = usgfw2.init_ultrasound_usgfw2() # Check probe if ERR == 2: logger.error("Main Usgfw2 library object not created") usgfw2.Close_and_release() sys.exit() ERR = usgfw2.find_connected_probe() if ERR != 101: logger.error("Probe not detected") usgfw2.Close_and_release() sys.exit() ERR = usgfw2.data_view_function() if ERR < 0: logger.error( "Main ultrasound scanning object for selected probe not created" ) sys.exit() ERR = usgfw2.mixer_control_function(0, 0, w, h, 0, 0, 0) if ERR < 0: logger.error("B mixer control not returned") sys.exit() # Probe setting res_X = ctypes.c_float(0.0) res_Y = ctypes.c_float(0.0) usgfw2.get_resolution(ctypes.pointer(res_X), ctypes.pointer(res_Y)) X_axis = np.zeros(shape=(w)) Y_axis = np.zeros(shape=(h)) if w % 2 == 0: k = 0 for i in range(-w // 2, w // 2 + 1): if i < 0: j = i + 0.5 X_axis[k] = j * res_X.value k = k + 1 else: if i > 0: j = i - 0.5 X_axis[k] = j * res_X.value k = k + 1 else: for i in range(-w // 2, w // 2): X_axis[i + w / 2 + 1] = i * res_X.value for i in range(0, h - 1): Y_axis[i] = i * res_Y.value old_resolution_x = res_X.value old_resolution_y = res_X.value # Image setting p_array = (ctypes.c_uint * w * h * 4)() fig, ax = plt.subplots() usgfw2.return_pixel_values(ctypes.pointer(p_array)) buffer_as_numpy_array = np.frombuffer(p_array, np.uint) reshaped_array = np.reshape(buffer_as_numpy_array, (w, h, 4)) img = ax.imshow( reshaped_array[:, :, 0:3], cmap="gray", vmin=0, vmax=255, origin="lower", extent=[np.amin(X_axis), np.amax(X_axis), np.amax(Y_axis), np.amin(Y_axis)], ) # starting copy from the original __int__ self.w = w self.h = h ( self.usgfw2, self.p_array, self.res_X, self.res_Y, self.old_resolution_x, self.old_resolution_y, self.X_axis, self.Y_axis, self.img, ) = ( usgfw2, p_array, res_X, res_Y, old_resolution_x, old_resolution_y, X_axis, Y_axis, img, ) # return the image from telemed def imaging(self): self.usgfw2.return_pixel_values(ctypes.pointer(self.p_array)) buffer_as_numpy_array = np.frombuffer(self.p_array, np.uint) reshaped_array = np.reshape(buffer_as_numpy_array, (self.w, self.h, 4)) self.usgfw2.get_resolution( ctypes.pointer(self.res_X), ctypes.pointer(self.res_Y) ) if ( self.res_X.value != self.old_resolution_x or self.res_Y.value != self.old_resolution_y ): if self.w % 2 == 0: k = 0 for i in range(-self.w // 2, self.w // 2 + 1): if i < 0: j = i + 0.5 self.X_axis[k] = j * self.res_X.value k = k + 1 else: if i > 0: j = i - 0.5 self.X_axis[k] = j * self.res_X.value k = k + 1 else: for i in range(-self.w // 2, self.w // 2): self.X_axis[i + self.w / 2 + 1] = i * self.res_X.value for i in range(0, self.h - 1): self.Y_axis[i] = i * self.res_Y.value self.old_resolution_x = self.res_X.value self.old_resolution_y = self.res_X.value self.img.set_data(reshaped_array[:, :, 0:3]) self.img.set_extent( [ np.amin(self.X_axis), np.amax(self.X_axis), np.amax(self.Y_axis), np.amin(self.Y_axis), ] ) # Transfer image format to cv2 img_array = np.asarray(self.img.get_array()) img_array = img_array[::-1, :, ::-1] # format same as plt image, RBG to BGR return img_array class Thread(QThread): updateFrame = Signal(QImage) def __init__(self, parent=None, args=None): QThread.__init__(self, parent) self.status = True self.cap = True self.args = args # init telemed if args.video is None: self.telemed = Telemed() # init model is_async = ( True if self.args.jobs == "auto" or int(self.args.jobs) > 1 else False ) self.model = Model( model_path=self.args.model, imgsz=self.args.img_size, classes=self.args.classes, device=self.args.device, plot_mask=self.args.plot_mask, conf_thres=self.args.conf_thres, is_async=is_async, n_jobs=self.args.jobs, ) def get_stats_fig(self, aorta_widths, aorta_confs, fig_w, fig_h, ts): title_font_size = 28 body_font_size = 24 img_quality = 100 * np.mean(aorta_confs) avg_width = np.mean(aorta_widths) max_width = np.max(aorta_widths) suggestions = [ "N/A, within normal limit", "Follow up in 5 years", "Make an appointment as soon as possible", ] s = None if avg_width < 3: s = suggestions[0] elif avg_width < 5: s = suggestions[1] else: s = suggestions[2] # region smoothing: method 2, keep the peaks # peaks = signal.find_peaks(aorta_widths, height=0.5, distance=40) # new_y = [] # # smooth the values between the peaks # start = 0 # end = peaks[0][0] # new_y.extend(signal.savgol_filter(aorta_widths[start:end], end - start, 2)) # for i in range(len(peaks[0]) - 1): # start = peaks[0][i] + 1 # end = peaks[0][i + 1] # new_y.append(aorta_widths[peaks[0][i]]) # add peak value # new_y.extend( # signal.savgol_filter( # aorta_widths[start:end], # end - start, # window size used for filtering # 2, # ) # ) # order of fitted polynomial # # add the last peak # new_y.append(aorta_widths[peaks[0][-1]]) # start = peaks[0][-1] + 1 # end = len(aorta_widths) # new_y.extend(signal.savgol_filter(aorta_widths[start:end], end - start, 2)) # endregion # region smoothing: method 1, do not keep the peaks window_size = 53 if len(aorta_widths) < window_size: window_size = len(aorta_widths) - 1 new_y = signal.savgol_filter(aorta_widths, window_size, 3) # endregion x = np.arange(1, len(aorta_widths) + 1, dtype=int) fig = go.Figure() fig.add_trace( go.Scatter( x=x, y=aorta_widths, mode="lines", line=dict(color="royalblue", width=1) ) ) fig.add_trace( go.Scatter( x=x, y=new_y, mode="lines", marker=dict( size=3, color="mediumpurple", ), ) ) fig.update_layout( autosize=False, width=fig_w, height=fig_h, margin=dict(l=50, r=50, b=50, t=400, pad=4), paper_bgcolor="LightSteelBlue", showlegend=False, ) fig.add_annotation( text=f"max={max_width:.2f} cm", x=np.argmax(aorta_widths), y=np.max(aorta_widths), xref="x", yref="y", showarrow=True, font=dict(color="#ffffff"), arrowhead=2, arrowsize=1, arrowwidth=2, borderpad=4, bgcolor="#ff7f0e", opacity=0.8, ) fig.add_annotation( text=f"smoothed max={np.max(new_y):.2f} cm", x=np.argmax(new_y), y=np.max(new_y), xref="x", yref="y", showarrow=True, font=dict(color="#ffffff"), arrowhead=2, arrowsize=1, arrowwidth=2, ax=-100, ay=-50, borderpad=4, bgcolor="#ff7f0e", opacity=0.8, ) fig.add_annotation( text="Report of Abdominal Aorta Examination", xref="paper", yref="paper", x=0.5, y=2.3, showarrow=False, font=dict(size=title_font_size), ) fig.add_annotation( text=f"Image acquisition quality: {img_quality:.0f}%", xref="paper", yref="paper", x=0, y=2.0, showarrow=False, font=dict(size=body_font_size), ) fig.add_annotation( text=f"Aorta Maximal Width: {max_width:.2f} cm", xref="paper", yref="paper", x=0, y=1.8, showarrow=False, font=dict(size=body_font_size), ) fig.add_annotation( text=f"Aorta Maximal Width (Smoothed): {np.max(new_y):.2f} cm", xref="paper", yref="paper", x=0, y=1.6, showarrow=False, font=dict(size=body_font_size), ) fig.add_annotation( text=f"Average: {avg_width:.2f} cm", xref="paper", yref="paper", x=0, y=1.4, showarrow=False, font=dict(size=body_font_size), ) fig.add_annotation( text=f"Suggestion: {s}", xref="paper", yref="paper", x=0, y=1.2, showarrow=False, font=dict(size=body_font_size), ) fig.add_annotation( text=f"Generated at {ts}", xref="paper", yref="paper", x=1, y=1, showarrow=False, ) return fig def run(self): one_cm_in_pixels = 48 # hard-coded aorta_cm_thre1 = 3 aorta_cm_thre2 = 5 black = (0, 0, 0) white = (255, 255, 255) red = (0, 0, 255) green = (0, 255, 0) aorta_widths_stats = [0, 0, 0] # three ranges: <3, 3-5, >5 aorta_widths = [] aorta_confs = [] expected_fps = None frame_count = None frame_w = None frame_h = None if self.args.video: self.cap = cv2.VideoCapture(self.args.video) expected_fps = self.cap.get(cv2.CAP_PROP_FPS) secs_per_frame = 1 / expected_fps frame_w, frame_h = int(self.cap.get(cv2.CAP_PROP_FRAME_WIDTH)), int( self.cap.get(cv2.CAP_PROP_FRAME_HEIGHT) ) frame_count = int(self.cap.get(cv2.CAP_PROP_FRAME_COUNT)) logger.info(f"Video source FPS: {expected_fps}") logger.info(f"Milliseconds per frame: {secs_per_frame}") logger.info(f"Video source resolution (WxH): {frame_w}x{frame_h}") logger.info(f"Video source frame count: {frame_count}") assert frame_count > 0, "No frame found" n_read_frames = 0 next_frame_to_infer = 0 next_frame_to_show = 0 n_repeat_failure = 0 is_last_failed = False start_time = perf_counter() while self.status: frame = None # avoid infinite loop if n_repeat_failure > 30: break # inference color_frame, others, results, xyxy, conf = None, None, None, None, None if self.model.is_async: results = self.model.get_result(next_frame_to_show) if results: color_frame, others = results xyxy, conf, _ = others next_frame_to_show += 1 if self.model.is_async and self.model.is_free_to_infer_async(): if self.args.video: ret, frame = self.cap.read() if not ret: n_repeat_failure += 1 if is_last_failed else 0 is_last_failed = True continue else: # read the frame from telemed # TODO(martin): Check read failure frame = self.telemed.imaging() n_read_frames += 1 self.model.predict_async(frame, next_frame_to_infer) next_frame_to_infer += 1 elif not self.model.is_async: if self.args.video: ret, frame = self.cap.read() if not ret: n_repeat_failure += 1 if is_last_failed else 0 is_last_failed = True continue else: # read the frame from telemed # TODO(martin): Check read failure frame = self.telemed.imaging() n_read_frames += 1 results = self.model.predict(frame) color_frame, others = results xyxy, conf, _ = others # bbox and confidence if results is None: continue is_last_failed = False # check if aorta is within the ROI box, and draw the box aorta_width_in_cm = 0 is_found = xyxy is not None is_in_box = False is_too_left, is_too_right = False, False w, h = color_frame.shape[1], color_frame.shape[0] box_w = int(w * 0.1) box_h = int(h * 0.5) box_top_left = (w // 2 - box_w // 2, h // 4) box_bottom_right = (w // 2 + box_w // 2, h // 4 + box_h) if xyxy is not None: x1, y1, x2, y2 = xyxy # check aorta width aorta_width_in_cm = (x2 - x1) / one_cm_in_pixels aorta_widths.append(aorta_width_in_cm) aorta_confs.append(conf) if aorta_width_in_cm < aorta_cm_thre1: aorta_widths_stats[0] += 1 elif aorta_width_in_cm < aorta_cm_thre2: aorta_widths_stats[1] += 1 else: aorta_widths_stats[2] += 1 # check whether aorta is in the box if ( x1 > box_top_left[0] and x2 < box_bottom_right[0] and y1 > box_top_left[1] and y2 < box_bottom_right[1] ): is_in_box = True is_too_right = x2 > box_bottom_right[0] is_too_left = x1 < box_top_left[0] # plot ROI box with color status box_color = green if is_in_box else red color_frame = cv2.rectangle( color_frame, box_top_left, box_bottom_right, box_color, 2 ) assert not ( is_too_left and is_too_right ), "Cannot be both too left and too right" if is_too_left: start_p = (box_top_left[0], int(h * 0.9)) end_p = (box_bottom_right[0], int(h * 0.9)) cv2.arrowedLine(color_frame, start_p, end_p, red, 3) if is_too_right: start_p = (box_bottom_right[0], int(h * 0.9)) end_p = (box_top_left[0], int(h * 0.9)) cv2.arrowedLine(color_frame, start_p, end_p, red, 3) if is_in_box: cv2.putText( color_frame, "GOOD", (box_top_left[0], int(h * 0.9)), cv2.FONT_HERSHEY_SIMPLEX, 1, green, 3, ) # plot aorta width text = ( f"Aorta width: {aorta_width_in_cm:.2f} cm" if is_found else "Aorta width: N/A" ) cv2.putText( color_frame, text, (50, 90), cv2.FONT_HERSHEY_SIMPLEX, 1, white, 3 ) # region FPS fps = None if n_read_frames > 0: fps = n_read_frames / (perf_counter() - start_time) # Slow down the loop if FPS is too high if self.args.sync: while fps > expected_fps: time.sleep(0.001) fps = n_read_frames / (perf_counter() - start_time) cv2.putText( color_frame, f"FPS: {fps:.2f}", (50, 30), cv2.FONT_HERSHEY_SIMPLEX, 1, white, 3, ) # endregion # Creating and scaling QImage h, w, ch = color_frame.shape img = QImage(color_frame.data, w, h, ch * w, QImage.Format_BGR888) scaled_img = img.scaled(video_w, video_h, Qt.KeepAspectRatio) # Emit signal self.updateFrame.emit(scaled_img) if self.args.video: progress = 100 * n_read_frames / frame_count fps_msg = f", FPS: {fps:.2f}" if fps is not None else "" print( f"Processed {n_read_frames}/{frame_count} ({progress:.2f}%) frames" + fps_msg, end="\r" if n_read_frames < frame_count else os.linesep, ) if n_read_frames >= frame_count: logger.info("Finished processing video") break if self.args.video: self.cap.release() if not self.status: logger.info("Stopped by user") return # draw a black image with frame width & height # with some text in center indicating generating report # it's just a dummy step to make demo more real im = np.zeros((frame_h, frame_w, 3), np.uint8) cv2.putText( im, "Generating report for you...", (frame_w // 3, frame_h // 2), cv2.FONT_HERSHEY_SIMPLEX, 1, white, 3, ) img = QImage(im.data, frame_w, frame_h, ch * w, QImage.Format_BGR888) scaled_img = img.scaled(video_w, video_h, Qt.KeepAspectRatio) self.updateFrame.emit(scaled_img) time.sleep(3) # plot aorta width tracing line chart now_t = datetime.now() ts1 = now_t.strftime("%Y%m%d_%H%M%S") ts2 = now_t.strftime("%Y/%m/%d %I:%M:%S") Path("runs").mkdir(parents=True, exist_ok=True) # np.save("runs/aorta_widths.npy", aorta_widths) fig_out_p = f"runs/aorta_report_{ts1}.jpeg" fig = self.get_stats_fig(aorta_widths, aorta_confs, video_w, video_h, ts2) # This may hang under Windows: https://github.com/plotly/Kaleido/issues/110 # The workaround is to install older kaleido version (see requirements.txt) fig.write_image(fig_out_p) logger.info(f"Saved aorta report: {fig_out_p}") img_bytes = fig.to_image(format="jpg", width=video_w, height=video_h) line_chart = np.array(Image.open(io.BytesIO(img_bytes))) line_chart = cv2.cvtColor(line_chart, cv2.COLOR_RGB2BGR) h, w, ch = line_chart.shape img = QImage(line_chart.data, video_w, video_h, ch * w, QImage.Format_BGR888) scaled_img = img.scaled(w, h, Qt.KeepAspectRatio) # Emit signal self.updateFrame.emit(scaled_img) time.sleep(5) # keep report open until user closes the window while self.status and not self.args.exit_on_end: time.sleep(0.1) class Window(QMainWindow): def __init__(self, args=None): super().__init__() # Title and dimensions self.setWindowTitle("Demo") self.setGeometry(0, 0, 800, 500) # Create a label for the display camera self.label = QLabel(self) # self.label.setFixedSize(self.width(), self.height()) self.label.setFixedSize(video_w, video_h) # Thread in charge of updating the image self.th = Thread(self, args) self.th.finished.connect(self.close) self.th.updateFrame.connect(self.setImage) # Buttons layout buttons_layout = QHBoxLayout() self.button1 = QPushButton("Start") self.button2 = QPushButton("Stop/Close") self.button1.setSizePolicy(QSizePolicy.Preferred, QSizePolicy.Expanding) self.button2.setSizePolicy(QSizePolicy.Preferred, QSizePolicy.Expanding) buttons_layout.addWidget(self.button2) buttons_layout.addWidget(self.button1) right_layout = QHBoxLayout() # right_layout.addWidget(self.group_model, 1) right_layout.addLayout(buttons_layout, 1) # Main layout layout = QVBoxLayout() layout.addWidget(self.label) layout.addLayout(right_layout) # Central widget widget = QWidget(self) widget.setLayout(layout) self.setCentralWidget(widget) # Connections self.button1.clicked.connect(self.start) self.button2.clicked.connect(self.kill_thread) self.button2.setEnabled(False) if args.start_on_open: # start thread self.start() @Slot() def kill_thread(self): logger.info("Finishing...") self.th.status = False time.sleep(1) # Give time for the thread to finish self.button2.setEnabled(False) self.button1.setEnabled(True) cv2.destroyAllWindows() self.th.exit() # Give time for the thread to finish time.sleep(1) @Slot() def start(self): logger.info("Starting...") self.button2.setEnabled(True) self.button1.setEnabled(False) self.th.start() logger.info("Thread started") @Slot(QImage) def setImage(self, image): self.label.setPixmap(QPixmap.fromImage(image)) if __name__ == "__main__": # get user inputs using argparse parser = argparse.ArgumentParser() parser.add_argument( "--video", type=str, default=None, help="path to video file, if None (default) would read from telemed", ) parser.add_argument( "--model", type=str, default="best_openvino_model/best.xml", help="path to model file", ) parser.add_argument("--img-size", type=int, default=640, help="image size") parser.add_argument( "--classes", nargs="+", type=int, default=[0], help="filter by class" ) parser.add_argument("--device", type=str, default="CPU", help="device to use") parser.add_argument("--sync", action="store_true", help="sync video FPS") parser.add_argument("--plot-mask", action="store_true", help="plot mask") parser.add_argument("--conf-thres", type=float, default=0.25, help="conf thresh") parser.add_argument("--jobs", type=str, default=1, help="num of jobs, async if > 1") parser.add_argument("--start-on-open", action="store_true", help="start on open") parser.add_argument("--exit-on-end", action="store_true", help="exit if video ends") args = parser.parse_args() assert ( args.jobs == "auto" or int(args.jobs) > 0 ), f"--jobs must be > 0 or auto, got {args.jobs}" if args.video: assert Path(args.video).exists(), f"Video file {args.video} not found" assert Path(args.model).exists(), f"Model file {args.model} not found" app = QApplication() w = Window(args) w.show() sys.exit(app.exec())