Merge pull request #1 from LazyBusyYang/init_commit

[Init] Init repo with README and stream helper main program

.github/workflows/lint.yml

@@ -0,0 +1,30 @@
+name: lint
+
+on: [push, pull_request]
+
+concurrency:
+  group: ${{ github.workflow }}-${{ github.ref }}
+  cancel-in-progress: true
+
+jobs:
+  lint:
+    runs-on: ubuntu-20.04
+    steps:
+      - uses: actions/checkout@v2
+      - name: Set up Python 3.8
+        uses: actions/setup-python@v2
+        with:
+          python-version: 3.8
+      - name: Install pre-commit hook
+        run: |
+          sudo apt-add-repository ppa:brightbox/ruby-ng -y
+          sudo apt-get update
+          sudo apt-get install -y ruby2.7
+          pip install pre-commit
+          pre-commit install
+      - name: Linting
+        run: pre-commit run --all-files
+      - name: Check docstring coverage
+        run: |
+          pip install interrogate
+          interrogate -vinmMI --ignore-init-method --ignore-module --ignore-nested-functions --ignore-regex "__repr__" -f 50 cat_stream/

.gitignore

@@ -0,0 +1,4 @@
+build
+yolov5s.pt
+*.egg-info*
+__pycache__

.isort.cfg

@@ -0,0 +1,2 @@
+[settings]
+known_third_party = cv2,numpy,setuptools,simpleobsws

.pre-commit-config.yaml

@@ -0,0 +1,45 @@
+exclude: .*/tests/data
+repos:
+  - repo: https://github.com/PyCQA/flake8.git
+    rev: 3.8.3
+    hooks:
+      - id: flake8
+  - repo: https://github.com/asottile/seed-isort-config.git
+    rev: v2.2.0
+    hooks:
+      - id: seed-isort-config
+        args: [--settings-path, ./]
+  - repo: https://github.com/PyCQA/isort.git
+    rev: 5.12.0
+    hooks:
+      - id: isort
+        args: [--settings-file, ./setup.cfg]
+  - repo: https://github.com/pre-commit/mirrors-yapf.git
+    rev: v0.30.0
+    hooks:
+      - id: yapf
+  - repo: https://github.com/pre-commit/pre-commit-hooks.git
+    rev: v3.1.0
+    hooks:
+      - id: trailing-whitespace
+        args: [--markdown-linebreak-ext=md]
+        exclude: .*/tests/data/
+      - id: check-yaml
+      - id: end-of-file-fixer
+      - id: requirements-txt-fixer
+      - id: double-quote-string-fixer
+      - id: check-merge-conflict
+      - id: fix-encoding-pragma
+        args: ["--remove"]
+      - id: mixed-line-ending
+        args: ["--fix=lf"]
+  - repo: https://github.com/codespell-project/codespell
+    rev: v2.1.0
+    hooks:
+      - id: codespell
+        args: ["--ignore-words-list", "ue"]
+  - repo: https://github.com/myint/docformatter.git
+    rev: v1.3.1
+    hooks:
+      - id: docformatter
+        args: ["--in-place", "--wrap-descriptions", "79"]

README.md

@@ -0,0 +1,45 @@
+# Cat Stream
+
+## Introduction
+
+This project is a multi-perspective cat live streaming program implemented based on Python and OBS Studio. It can select the camera view according to the detection results of cats from multiple RTSP network cameras in the home, control OBS to activate scenes with cats, and complete unattended cat live streaming on the network.
+
+![Rooms, cameras and cat](./resources/room_illustration.jpg)
+
+The main program initializes upon startup according to the configuration file, connects to the OBS websocket server, and enters a while loop. Within each iteration of the loop, it requests the latest RTSP URLs for each perspective from OBS. It then uses ffmpeg or cv2 to read the latest video frames from RTSP. Utilizing YOLOv5 or cv2, it detects whether a cat is present in the video frames. Based on the detection results, it sends scene-switching signals to OBS. The flowchart is shown in the figure below.
+
+![Mainloop flow chart](./resources/mainloop_flow.png)
+
+## Prerequisites
+
+Before you begin, ensure you have met the following requirements:
+
+- **OBS studio**: Please configure the various scenes and RTSP media sources in OBS, and enable OBS websocket. This project merely automates the process of scene switching, not creating.
+- **Pytorch**: Pytorch-CPU is required for YOLOv5 cat body detection. Without pytorch, the performance of cv2 cat face detection is relatively poor.
+- **FFmpeg command tool**: To read RTSP stream, ffmpeg works perfectly, while cv2 usually reports decoding errors(the program won't crash).
+
+## Installing
+
+Navigate to the root directory of the project, and then directly use pip to install this project. Third-party dependencies will be automatically completed.
+```bash
+pip install .
+```
+To run the YOLO object detection with GPU, please refer to the official PyTorch tutorial and modify the device-related section in the configuration file of this project.
+
+## Running
+
+Please write the configuration file needed for live streaming control based on the existing configuration files and runtime environment in the `configs/` directory. For the specific meaning of each value in the configuration file, you can refer to the docstring of the `__init__` function in the code.
+
+When the configuration file is ready, start program with a command like below:
+```bash
+python tools/main.py --config_path configs/yolov5_ffmpeg_3scenes.py
+```
+
+
+## Authors
+
+* **LazyBusyYang** - [Github Page](https://github.com/LazyBusyYang)
+
+## License
+
+This project is licensed under the Apache 2.0 License

cat_stream/config_reader.py

@@ -0,0 +1,27 @@
+import importlib
+import os
+import sys
+import types
+
+
+def file2dict(file_path: str) -> dict:
+    """Convert a python file to a dict.
+
+    Args:
+        file_path (str): The path of the file.
+    Returns:
+        dict: The dict converted from the file.
+    """
+    file_name = os.path.basename(file_path)
+    file_name = file_name.split('.')[0]
+    file_dir = os.path.dirname(file_path)
+    sys.path.insert(0, file_dir)
+    file_module = importlib.import_module(file_name)
+    sys.path.pop(0)
+    file_dict = {
+        name: value
+        for name, value in file_module.__dict__.items() if
+        not name.startswith('__') and not isinstance(value, types.ModuleType)
+        and not isinstance(value, types.FunctionType)
+    }
+    return file_dict

cat_stream/detection.py

@@ -0,0 +1,78 @@
+# import Abstract class
+import cv2
+import numpy as np
+from abc import ABC
+from typing import Any
+
+try:
+    import torch
+    has_torch = True
+except ImportError:
+    has_torch = False
+
+
+class BaseCatDetection(ABC):
+
+    def __init__(self) -> None:
+        pass
+
+    def detect(self, frame: np.ndarray) -> Any:
+        pass
+
+    def check_cat(self, detect_result: Any) -> bool:
+        pass
+
+
+class OpenCVCatFaceDetection(BaseCatDetection):
+    """A cat face detection class using OpenCV, only for poor machines as cat
+    butt detection is not supported."""
+
+    def __init__(self) -> None:
+        super().__init__()
+
+    def detect(self, frame: np.ndarray) -> list:
+        cat_cascade = cv2.CascadeClassifier(cv2.data.haarcascades +
+                                            'haarcascade_frontalcatface.xml')
+        gray_image = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
+        # TODO: tune the parameters if needed
+        cat_faces = cat_cascade.detectMultiScale(
+            gray_image, scaleFactor=1.1, minNeighbors=5, minSize=(30, 30))
+        return cat_faces
+
+    def check_cat(self, detect_result: Any) -> bool:
+        if len(detect_result) > 0:
+            return True
+        else:
+            return False
+
+
+class YOLOv5CatDetection(BaseCatDetection):
+    """A cat detection class using YOLOv5."""
+
+    def __init__(self, device: str = 'cpu') -> None:
+        super().__init__()
+        if not has_torch:
+            raise ImportError('PyTorch is not installed.')
+        # TODO: Load locally if needed
+        self.device = device
+        self.model = torch.hub.load(
+            'ultralytics/yolov5', 'yolov5s', pretrained=True).to(self.device)
+
+    def detect(self, frame: np.ndarray) -> Any:
+        results = self.model(frame)
+        return results
+
+    def check_cat(self, detect_result: Any) -> bool:
+        return 'cat' in str(detect_result)
+
+
+def build_detection(cfg: dict) -> BaseCatDetection:
+    cfg = cfg.copy()
+    class_name = cfg.pop('type')
+    if class_name == 'OpenCVCatFaceDetection':
+        detection = OpenCVCatFaceDetection(**cfg)
+    elif class_name == 'YOLOv5CatDetection':
+        detection = YOLOv5CatDetection(**cfg)
+    else:
+        raise TypeError(f'Invalid type {class_name}.')
+    return detection