loopbio blog

  """OpenCV utils."""
  import contextlib
  import os
  import cv2
  import logging

  _log = logging.getLogger(__package__)


  @contextlib.contextmanager
  def cv2_num_threads(num_threads):
      """Context manager to temporarily change the number of threads used by opencv."""
      old_num_threads = cv2.getNumThreads()
      cv2.setNumThreads(num_threads)
      yield
      cv2.setNumThreads(old_num_threads)


  # A string to request not to change the current value of an envvar
  USE_CURRENT_VALUE = object()


  @contextlib.contextmanager
  def envvar(name, value=USE_CURRENT_VALUE):
      """
      Context manager to temporarily change the value of an environment variable for the current process.

      Remember that some envvars only affects the process on startup (e.g. LD_LIBRARY_PATH).

      Parameters
      ----------
      name : string
        The name of the environment value to modify.

      value : None, `cv2utils.USE_CURRENT_VALUE` or object; default "USE_CURRENT_VALUE"
        If `cv2utils.USE_CURRENT_VALUE`, the environment variable value is not modified whatsoever.
        If None, the environment variable value is temporarily removed, if it exists.
        Else, str(value) will be temporarily set as the value for the environment variable

      Examples
      --------
      When a variable is not already set...
      >>> name = 'AN_ENVIRONMENT_VARIABLE'
      >>> with envvar(name, None):
      ...     print(os.environ.get(name))
      None
      >>> with envvar(name, USE_CURRENT_VALUE):
      ...     print(os.environ.get(name))
      None
      >>> with envvar(name, 42):
      ...     print(os.environ.get(name))
      42
      >>> print(os.environ.get(name))
      None

      When a variable is already set...
      >>> os.environ[name] = 'a_default_value'
      >>> with envvar(name, USE_CURRENT_VALUE):
      ...     print(os.environ.get(name))
      a_default_value
      >>> with envvar(name, None):
      ...     print(os.environ.get(name))
      None
      >>> print(os.environ.get(name))
      a_default_value
      >>> with envvar(name, 42):
      ...     print(os.environ.get(name))
      42
      >>> print(os.environ.get(name))
      a_default_value
      """
      if value is USE_CURRENT_VALUE:
          yield
      elif name not in os.environ:
          if value is not None:
              os.environ[name] = str(value)
              yield
              del os.environ[name]
          else:
              yield
      else:
          old_value = os.environ[name]
          if value is not None:
              os.environ[name] = str(value)
          else:
              del os.environ[name]
          yield
          os.environ[name] = old_value


  def ffmpeg_thread_count(thread_count=USE_CURRENT_VALUE):
      """
      Context manager to temporarily change the number of threads requested by cv2.VideoCapture.

      This works manipulating global state, so this function is not thread safe. Take care
      if you instantiate capture objects with different thread_count concurrently.

      The actual behavior depends on the codec. Some codecs will honor thread_count,
      while others will not. You can always call `video_capture_thread_count(cap)`
      to check whether the concrete codec used does one thing or the other.

      Note that as of 2018/03, we only support changing the number of threads for decoding
      (i.e. VideoCapture, but not VideoWriter).

      Parameters
      ----------
      thread_count : int or None or `cv2utils.USE_CURRENT_VALUE`, default USE

        * if None, then no change on the default behavior of opencv will happen
          on opencv 3.4.1 and linux, this means "the number of logical cores as reported
          by "sysconf(SC_NPROCESSORS_ONLN)" - which is a pretty aggresive setting in terms
          of resource consumption, specially in multiprocess applications,
          and might even be problematic if running with capped resources,
          like in a cgroups/container, under tasksel or numactl.

        * if an integer, set capture decoders to the specifiednumber of threads
          usually 0 means "auto", that is, let ffmpeg decide

        * if `cv2utils.USE_CURRENT_VALUE`, the current value of the environment
          variable OPENCV_FFMPEG_THREAD_COUNT is used (if undefined, then the default
          value given by opencv is used)
      """
      return envvar(name='OPENCV_FFMPEG_THREAD_COUNT', value=thread_count)


  def cv2_supports_thread_count():
      """Returns True iff opencv has been built with support to expose ffmpeg thread_count."""
      return hasattr(cv2, 'CAP_PROP_THREAD_COUNT')


  def video_capture_thread_count(cap):
      """
      Returns the number of threads used by a VideoCapture as reported by opencv.
      Returns None if the opencv build does not support this feature.
      """
      try:
          # noinspection PyUnresolvedReferences
          return cap.get(cv2.CAP_PROP_THREAD_COUNT)
      except AttributeError:
          return None


  def open_video_capture(path,
                         num_threads=USE_CURRENT_VALUE,
                         fail_if_unsupported_num_threads=False,
                         backend=cv2.CAP_FFMPEG):
      """
      Returns a VideoCapture object for the specified path.

      Parameters
      ----------
      path : string
        The path to a video source (file or device)

      num_threads : None, int or `cv2utils.USE_CURRENT_VALUE`, default None
        The number of threads used for decoding.
        If None, opencv defaults is used (number of logical cores in the system).
        If an int, the number of threads to use. Usually 0 means "auto", 1 "single-threaded"
        (but it might depend on the codec).

      fail_if_unsupported_num_threads : bool, default False
        If False, an warning is cast if num_threads is not None and setting the
        number of threads is unsupported either by opencv or the used codec.

        If True, a ValueError is raised in any of these two cases.

      backend : cv2 backend or None, default cv2.CAP_FFMPEG
        If provided, it will be used as preferred backend for opencv VidecCapture
      """
      if num_threads is not None and not cv2_supports_thread_count():
          message = ('OpenCV does not support setting the number of threads to %r; '
                     'use loopbio build' % num_threads)
          if fail_if_unsupported_num_threads:
              raise ValueError(message)
          else:
              _log.warn(message)

      with ffmpeg_thread_count(num_threads):
          if backend is not None:
              cap = cv2.VideoCapture(path, backend)
          else:
              cap = cv2.VideoCapture(path)

      if cap is None or not cap.isOpened():
          raise IOError("OpenCV unable to open %s" % path)

      if num_threads is USE_CURRENT_VALUE:
          try:
              num_threads = float(os.environ['OPENCV_FFMPEG_THREAD_COUNT'])
          except (KeyError, TypeError):
              num_threads = None
      if num_threads is not None and num_threads != video_capture_thread_count(cap):
          message = 'OpenCV num_threads for decoder setting to %r ignored for %s' % (num_threads, path)
          if fail_if_unsupported_num_threads:
              raise ValueError(message)
          else:
              _log.warn(message)

      return cap

  # Do whatever you need
  if not cap.isOpened():
      raise Exception('Something is wrong and the capture is not open')
  retval, image = cap.read()