Merge pull request #577 from luxonis/add_download_progress

Add download progress indicator in Qt GUI

depthai_demo.py

@@ -76,7 +76,7 @@ def run_all(self, conf):
         self.setup(conf)
         self.run()
 
-    def __init__(self, displayFrames=True, onNewFrame = noop, onShowFrame = noop, onNn = noop, onReport = noop, onSetup = noop, onTeardown = noop, onIter = noop, shouldRun = lambda: True, collectMetrics=False):
+    def __init__(self, displayFrames=True, onNewFrame = noop, onShowFrame = noop, onNn = noop, onReport = noop, onSetup = noop, onTeardown = noop, onIter = noop, shouldRun = lambda: True, showDownloadProgress=None, collectMetrics=False):
         self._openvinoVersion = None
         self._displayFrames = displayFrames
         self.toggleMetrics(collectMetrics)
@@ -89,8 +89,9 @@ def __init__(self, displayFrames=True, onNewFrame = noop, onShowFrame = noop, on
         self.onTeardown = onTeardown
         self.onIter = onIter
         self.shouldRun = shouldRun
-
-    def setCallbacks(self, onNewFrame=None, onShowFrame=None, onNn=None, onReport=None, onSetup=None, onTeardown=None, onIter=None, shouldRun=None):
+        self.showDownloadProgress = showDownloadProgress
+
+    def setCallbacks(self, onNewFrame=None, onShowFrame=None, onNn=None, onReport=None, onSetup=None, onTeardown=None, onIter=None, shouldRun=None, showDownloadProgress=None):
         if onNewFrame is not None:
             self.onNewFrame = onNewFrame
         if onShowFrame is not None:
@@ -107,6 +108,8 @@ def setCallbacks(self, onNewFrame=None, onShowFrame=None, onNn=None, onReport=No
             self.onIter = onIter
         if shouldRun is not None:
             self.shouldRun = shouldRun
+        if showDownloadProgress is not None:
+            self.showDownloadProgress = showDownloadProgress
 
     def toggleMetrics(self, enabled):
         if enabled:
@@ -136,6 +139,7 @@ def setup(self, conf: ConfigManager):
             self._blobManager = BlobManager(
                 zooDir=DEPTHAI_ZOO,
                 zooName=self._conf.getModelName(),
+                progressFunc=self.showDownloadProgress
             )
             self._nnManager = NNetManager(inputSize=self._conf.inputSize)
 
@@ -510,6 +514,7 @@ def runQt():
 
     class WorkerSignals(QObject):
         updateConfSignal = pyqtSignal(list)
+        updateDownloadProgressSignal = pyqtSignal(int, int)
         updatePreviewSignal = pyqtSignal(np.ndarray)
         setDataSignal = pyqtSignal(list)
         exitSignal = pyqtSignal()
@@ -538,7 +543,7 @@ def __init__(self, instance, parent, conf, selectedPreview=None):
         def run(self):
             self.running = True
             self.signals.setDataSignal.emit(["restartRequired", False])
-            self.instance.setCallbacks(shouldRun=self.shouldRun, onShowFrame=self.onShowFrame, onSetup=self.onSetup)
+            self.instance.setCallbacks(shouldRun=self.shouldRun, onShowFrame=self.onShowFrame, onSetup=self.onSetup, showDownloadProgress=self.showDownloadProgress)
             self.conf.args.bandwidth = "auto"
             if self.conf.args.deviceId is None:
                 devices = dai.Device.getAllAvailableDevices()
@@ -594,6 +599,9 @@ def onShowFrame(self, frame, source):
             if source == self.selectedPreview:
                 self.signals.updatePreviewSignal.emit(frame)
 
+        def showDownloadProgress(self, curr, total):
+            self.signals.updateDownloadProgressSignal.emit(curr, total)
+
         def onSetup(self, instance):
             if "onSetup" in self.file_callbacks:
                 self.file_callbacks["onSetup"](instance)
@@ -652,6 +660,7 @@ def start(self):
             self.running = True
             self.worker = Worker(self._demoInstance, parent=self, conf=self.confManager, selectedPreview=self.selectedPreview)
             self.worker.signals.updatePreviewSignal.connect(self.updatePreview)
+            self.worker.signals.updateDownloadProgressSignal.connect(self.updateDownloadProgress)
             self.worker.signals.setDataSignal.connect(self.setData)
             self.worker.signals.errorSignal.connect(self.showError)
             self.threadpool.start(self.worker)

depthai_sdk/requirements.txt

@@ -2,6 +2,6 @@ numpy>=1.19; python_version < "3.7"
 numpy>=1.21; python_version >= "3.7"
 opencv-python>4
 opencv-contrib-python>4
-blobconverter>=1.2.6
+blobconverter>=1.2.8
 pytube>=11.0.1
 depthai>2

depthai_sdk/src/depthai_sdk/managers/blob_manager.py

@@ -19,14 +19,18 @@ class BlobManager:
     _useBlob = False
     _zooModels = []
 
-    def __init__(self, blobPath:Path=None, configPath:Path=None, zooName:str=None, zooDir:Path=None):
+    def __init__(self, blobPath:Path=None, configPath:Path=None, zooName:str=None, zooDir:Path=None, progressFunc=None):
         """
         Args:
             blobPath (pathlib.Path, Optional): Path to the compiled MyriadX blob file
             configPath (pathlib.Path, Optional): Path to model config file that is used to download the model
             zooName (str, Optional): Model name to be taken from model zoo
             zooDir (pathlib.Path, Optional): Path to model zoo directory
+            progressFunc (func, Optional): Custom method to show download progress, should accept two arguments - current bytes and max bytes.
         """
+        if progressFunc is not None:
+            blobconverter.set_defaults(progress_func=progressFunc)
+
         if blobPath is not None:
             self._blobPath = blobPath
             self._useBlob = True
@@ -54,7 +58,7 @@ def __init__(self, blobPath:Path=None, configPath:Path=None, zooName:str=None, z
             self._configPath = configPath
 
 
-    def getBlob(self, shaves:int, openvinoVersion: dai.OpenVINO.Version):
+    def getBlob(self, shaves:int, openvinoVersion: dai.OpenVINO.Version, zooType:str = None):
         """
         This function is responsible for returning a ready to use MyriadX blob once requested.
         It will compile the model automatically using our online blobconverter tool. The compilation process will be
@@ -63,6 +67,7 @@ def getBlob(self, shaves:int, openvinoVersion: dai.OpenVINO.Version):
         Args:
             shaves (int): Specify how many shaves the model will use. Range 1-16
             openvinoVersion (depthai.OpenVINO.Version): OpenVINO version which will be used to compile the MyriadX blob
+            zooType (str, Optional): Specifies model zoo type to download blob from
 
         Returns:
             pathlib.Path: Path to compiled MyriadX blob
@@ -80,7 +85,8 @@ def getBlob(self, shaves:int, openvinoVersion: dai.OpenVINO.Version):
                 self._blobPath = blobconverter.from_zoo(
                     name=self._zooName,
                     shaves=shaves,
-                    version=version
+                    version=version,
+                    zoo_type=zooType
                 )
                 self._useBlob = True
                 return self._blobPath

depthai_sdk/src/depthai_sdk/utils.py

@@ -226,6 +226,8 @@ def cropToAspectRatio(frame, size):
     Args:
         frame (numpy.ndarray): Source frame that will be cropped
         size (tuple): Desired frame size (width, heigth)
+    Returns:
+         numpy.ndarray: Cropped frame
     """
     shape = frame.shape
     h = shape[0]
@@ -252,6 +254,8 @@ def resizeLetterbox(frame, size):
     Args:
         frame (numpy.ndarray): Source frame that will be resized
         size (tuple): Desired frame size (width, heigth)
+    Returns:
+         numpy.ndarray: Resized frame
     """
     border_v = 0
     border_h = 0
@@ -262,3 +266,25 @@ def resizeLetterbox(frame, size):
     frame = cv2.copyMakeBorder(frame, border_v, border_v, border_h, border_h, cv2.BORDER_CONSTANT, 0)
     return cv2.resize(frame, size)
 
+
+def createBlankFrame(width, height, rgb_color=(0, 0, 0)):
+    """
+    Create new image(numpy array) filled with certain color in RGB
+
+    Args:
+        width (int): New frame width
+        height (int): New frame height
+        rgb_color (tuple, Optional): Specify frame fill color in RGB format (default (0,0,0) - black)
+
+    Returns:
+         numpy.ndarray: New frame filled with specified color
+    """
+    # Create black blank image
+    image = np.zeros((height, width, 3), np.uint8)
+
+    # Since OpenCV uses BGR, convert the color first
+    color = tuple(reversed(rgb_color))
+    # Fill image with color
+    image[:] = color
+
+    return image

gui/main.py

@@ -2,6 +2,8 @@
 import sys
 from pathlib import Path
 
+import blobconverter
+import cv2
 from PyQt5.QtQml import QQmlApplicationEngine, qmlRegisterType, qmlRegisterSingletonType, QQmlEngine
 from PyQt5.QtQuick import QQuickPaintedItem
 from PyQt5.QtGui import QImage
@@ -11,7 +13,7 @@
 # To be used on the @QmlElement decorator
 # (QML_IMPORT_MINOR_VERSION is optional)
 from PyQt5.QtWidgets import QApplication
-from depthai_sdk import Previews, resizeLetterbox
+from depthai_sdk import Previews, resizeLetterbox, createBlankFrame
 
 
 class Singleton(type(QQuickPaintedItem)):
@@ -266,6 +268,7 @@ class DemoQtGui:
     instance = None
     writer = None
     window = None
+    progressFrame = None
 
     def __init__(self):
         global instance
@@ -299,8 +302,42 @@ def updatePreview(self, frame):
             img = QImage(scaledFrame.data, w, h, w, 24)  # 24 - QImage.Format_Grayscale8
         self.writer.update_frame(img)
 
+    def updateDownloadProgress(self, curr, total):
+        frame = self.createProgressFrame(curr / total)
+        img = QImage(frame.data, frame.shape[1], frame.shape[0], frame.shape[2] * frame.shape[1], 29)  # 29 - QImage.Format_BGR888
+        self.writer.update_frame(img)
+
+    def createProgressFrame(self, donePercentage=None):
+        confManager = getattr(self, "confManager", None)
+        w, h = int(self.writer.width()), int(self.writer.height())
+        if self.progressFrame is None:
+            self.progressFrame = createBlankFrame(w, h)
+            if confManager is None:
+                downloadText = "Downloading model blob..."
+            else:
+                downloadText = f"Downloading {confManager.getModelName()} blob..."
+            textsize = cv2.getTextSize(downloadText, cv2.FONT_HERSHEY_TRIPLEX, 0.5, 4)[0][0]
+            offset = int((w - textsize) / 2)
+            cv2.putText(self.progressFrame, downloadText, (offset, 250), cv2.FONT_HERSHEY_TRIPLEX, 0.5, (255, 255, 255), 4, cv2.LINE_AA)
+            cv2.putText(self.progressFrame, downloadText, (offset, 250), cv2.FONT_HERSHEY_TRIPLEX, 0.5, (0, 0, 0), 1, cv2.LINE_AA)
+
+        newFrame = self.progressFrame.copy()
+        if donePercentage is not None:
+            cv2.rectangle(newFrame, (100, 300), (460, 350), (255, 255, 255), cv2.FILLED)
+            cv2.rectangle(newFrame, (110, 310), (int(110 + 340 * donePercentage), 340), (0, 0, 0), cv2.FILLED)
+        return newFrame
+
+    def showSetupFrame(self, text):
+        w, h = int(self.writer.width()), int(self.writer.height())
+        setupFrame = createBlankFrame(w, h)
+        cv2.putText(setupFrame, text, (200, 250), cv2.FONT_HERSHEY_TRIPLEX, 0.5, (255, 255, 255), 4, cv2.LINE_AA)
+        cv2.putText(setupFrame, text, (200, 250), cv2.FONT_HERSHEY_TRIPLEX, 0.5, (0, 0, 0), 1, cv2.LINE_AA)
+        img = QImage(setupFrame.data, w, h, setupFrame.shape[2] * w, 29)  # 29 - QImage.Format_BGR888
+        self.writer.update_frame(img)
+
     def startGui(self):
         self.writer = self.window.findChild(QObject, "writer")
+        self.showSetupFrame("Starting demo...")
         medianChoices = list(filter(lambda name: name.startswith('KERNEL_') or name.startswith('MEDIAN_'), vars(dai.MedianFilter).keys()))[::-1]
         self.setData(["medianChoices", medianChoices])
         colorChoices = list(filter(lambda name: name[0].isupper(), vars(dai.ColorCameraProperties.SensorResolution).keys()))
@@ -310,4 +347,5 @@ def startGui(self):
         self.setData(["modelSourceChoices", [Previews.color.name, Previews.left.name, Previews.right.name]])
         versionChoices = sorted(filter(lambda name: name.startswith("VERSION_"), vars(dai.OpenVINO).keys()), reverse=True)
         self.setData(["ovVersions", versionChoices])
+        self.createProgressFrame()
         return self.app.exec()