Supress TensorFlow warnings; update dependencies

Also use PrettyTable to list source resolution

.python-version

@@ -1 +1 @@
-3.10.5
+3.11.5

README.md

@@ -49,6 +49,7 @@ This assumes you have Python 3.10 or 3.11 installed
 
 #### Poetry  
 1. `poetry install`  
+    a. For GPU support, use `poetry install -E cuda --with gpu`
 2. `poetry run -- wyzely-detect`  
 ### Configuration  
 The following are some basic CLI options. Most flags have environment variable equivalents which can be helpful when using Docker. 

pyproject.toml

@@ -21,11 +21,12 @@ ultralytics = "^8.0.190"
 hjson = "^3.1.0"
 numpy = "^1.23.2"
 
-# https://github.com/python-poetry/poetry/issues/6409
+# https://github.com/python-poetry/poetry/issues/6409#issuecomment-1911735833
-torch = ">=2.0.0, !=2.0.1, !=2.1.0"
+# To install with GPU, use poetry install -E cuda --with gpu
+torch = {version = "2.1.*", source = "pytorch-cpu", markers = "extra!='cuda'" }
 
 # https://stackoverflow.com/a/76477590/18270659
-# https://discuss.tensorflow.org/t/tensorflow-io-gcs-filesystem-with-windows/18849/4
+# https://discfuss.tensorflow.org/t/tensorflow-io-gcs-filesystem-with-windows/18849/4
 # Might be able to remove this version constraint later
 # Working versions:
 # Python version 3.10.12 and 3.10.5 both work
@@ -33,10 +34,33 @@ torch = ">=2.0.0, !=2.0.1, !=2.1.0"
 # cuDNN version - 8.8.1
 # Installed from Nvidia website - nvidia-cuda-toolkit is not installed, but default PopOS drivers are installed
 tensorflow-io-gcs-filesystem = "0.31.0"
-tensorflow = {version = "^2.14.0", extras = ["and-cuda"]}
+tensorflow = {version = "^2.14.0", markers = "extra!='cuda'"}
 
 
 deepface = "^0.0.79"
+prettytable = "^3.9.0"
+
+
+[tool.poetry.group.gpu]
+optional = true
+
+[tool.poetry.group.gpu.dependencies]
+torch = {version = "2.1.*", source = "pytorch-cu121", markers = "extra=='cuda'"}
+tensorflow = {version = "^2.14.0", extras = ["and-cuda"], markers = "extra=='cuda'"}
+
+[tool.poetry.extras]
+# Might be better to rename this to nocpu since it's more accurate
+cuda = []
+
+[[tool.poetry.source]]
+name = "pytorch-cpu"
+url = "https://download.pytorch.org/whl/cpu"
+priority = "explicit"
+
+[[tool.poetry.source]]
+name = "pytorch-cu121"
+url = "https://download.pytorch.org/whl/cu121"
+priority = "explicit"
 
 [tool.poetry.group.dev.dependencies]
 black = "^23.9.1"

wyzely_detect/__main__.py

@@ -1,8 +1,9 @@
 # import face_recognition
 from pathlib import Path
 import cv2
+import os
 
-from PrettyTable import PrettyTable
+from prettytable import PrettyTable
 
 # import hjson as json
 import torch
@@ -18,7 +19,7 @@ args = None
 def main():
     global objects_and_peoples
     global args 
-    # RUN_BY_COMPOSE = os.getenv("RUN_BY_COMPOSE") # Replace this with code to check for gpu
+
 
     args = argparser.parse_args()
 
@@ -49,7 +50,9 @@ def main():
     # Set the video capture to the appropriate source
     if not args.rtsp_url and not args.capture_device:
         print("No stream or capture device set, defaulting to capture device 0")
-        args.capture_device = [0]
+        video_sources = {
+            "devices": [cv2.VideoCapture(0)]
+        }
     else:
         video_sources = {
             "streams": [cv2.VideoCapture(url) for url in args.rtsp_url],
@@ -60,13 +63,22 @@ def main():
     # This makes it so that the video capture will only grab the most recent frame
     # However, this means that the video may be choppy
     # Only do this for streams
-    for stream in video_sources["streams"]:
+    try:
-        stream.set(cv2.CAP_PROP_BUFFERSIZE, 1)
+        for stream in video_sources["streams"]:
+            stream.set(cv2.CAP_PROP_BUFFERSIZE, 1)
+    # If there are no streams, this will throw a KeyError
+    except KeyError:
+        pass
+
+    # Print out the resolution of the video sources. Ideally, change this so the device ID/url is also printed
+    pretty_table = PrettyTable(field_names=["Source Type", "Resolution"])
+    for source_type, sources in video_sources.items():
+        for source in sources:
+            pretty_table.add_row(
+                [source_type, f"{source.get(cv2.CAP_PROP_FRAME_WIDTH)}x{source.get(cv2.CAP_PROP_FRAME_HEIGHT)}"]
+            )
+    print(pretty_table)
 
-    # Print the resolution of the video
-    print(
-        f"Video resolution: {video_capture.get(cv2.CAP_PROP_FRAME_WIDTH)}x{video_capture.get(cv2.CAP_PROP_FRAME_HEIGHT)}"  # noqa: E501
-    )
     print
     print("Beginning video capture...")
     while True:

wyzely_detect/utils/utils.py

@@ -1,8 +1,13 @@
 import cv2
+import os
 import numpy as np
 from pathlib import Path
-from deepface import DeepFace
+# https://stackoverflow.com/a/42121886/18270659
-from . import notify
+os.environ['TF_CPP_MIN_LOG_LEVEL']='3'
+
+
+from deepface import DeepFace # noqa: E402
+from . import notify # noqa: E402
 
 first_face_try = True