Feature/switch to detection model (#77)

* load model from ultralytics

* remove pyrovision deps

* update docker file

* set default weights

* comment

* use onnx

* fix deps

* fix test

* move to python 3.8

* wrong deletion

* extract pred

* missing numpy

* python 3.8

* remove install via req

* fix docs deps

.github/workflows/doc-deploy.yml

@@ -10,10 +10,10 @@ jobs:
       - uses: actions/checkout@v2
         with:
           persist-credentials: false
-      - name: Set up Python 3.6
+      - name: Set up Python 3.8
         uses: actions/setup-python@v1
         with:
-          python-version: 3.6
+          python-version: 3.8
           architecture: x64
       - name: Cache python modules
         uses: actions/cache@v2

.github/workflows/gh-page.yml

@@ -6,10 +6,10 @@ jobs:
   see-page-build-payload:
     runs-on: ubuntu-latest
     steps:
-      - name: Set up Python 3.6
+      - name: Set up Python 3.8
         uses: actions/setup-python@v1
         with:
-          python-version: 3.6
+          python-version: 3.8
           architecture: x64
       - name: check status
         run: |

.github/workflows/main.yml

@@ -11,10 +11,10 @@ jobs:
     runs-on: ubuntu-latest
     steps:
       - uses: actions/checkout@v2
-      - name: Set up Python 3.6
+      - name: Set up Python 3.8
         uses: actions/setup-python@v1
         with:
-          python-version: 3.6
+          python-version: 3.8
           architecture: x64
       - name: Cache python modules
         uses: actions/cache@v2
@@ -35,10 +35,10 @@ jobs:
     needs: pkg-install
     steps:
       - uses: actions/checkout@v2
-      - name: Set up Python 3.6
+      - name: Set up Python 3.8
         uses: actions/setup-python@v1
         with:
-          python-version: 3.6
+          python-version: 3.8
           architecture: x64
       - name: Cache python modules
         uses: actions/cache@v2
@@ -71,10 +71,10 @@ jobs:
     runs-on: ubuntu-latest
     steps:
       - uses: actions/checkout@v2
-      - name: Set up Python 3.6
+      - name: Set up Python 3.8
         uses: actions/setup-python@v1
         with:
-          python-version: 3.6
+          python-version: 3.8
           architecture: x64
       - name: Run flake8
         run: |
@@ -87,10 +87,10 @@ jobs:
     needs: pkg-install
     steps:
       - uses: actions/checkout@v2
-      - name: Set up Python 3.6
+      - name: Set up Python 3.8
         uses: actions/setup-python@v1
         with:
-          python-version: 3.6
+          python-version: 3.8
           architecture: x64
       - name: Cache python modules
         uses: actions/cache@v2

Dockerfile

@@ -1,18 +1,20 @@
-FROM pyronear/pyro-vision:latest
+FROM python:3.8.1-slim
 
 ENV PYTHONUNBUFFERED 1
 ENV PYTHONDONTWRITEBYTECODE 1
 
 
-COPY ./requirements.txt /tmp/requirements.txt
 COPY ./README.md /tmp/README.md
 COPY ./setup.py /tmp/setup.py
 COPY ./pyroengine /tmp/pyroengine
+COPY ./requirements.txt /tmp/requirements.txt
 
-RUN pip install --upgrade pip setuptools wheel \
+RUN apt update \
+    && apt install -y git \
+    && apt install ffmpeg libsm6 libxext6  -y \
+    && apt install -y gcc python3-dev \
+    && apt install --upgrade pip setuptools wheel \
     && pip install -e /tmp/. \
     && pip cache purge \
-    && apt-get autoremove -y \
-    && rm -rf /var/lib/apt/lists/* \
     && rm -rf /root/.cache/pip
 

docs/requirements.txt

@@ -1,2 +1,3 @@
-sphinx<3.5.0
 sphinx-rtd-theme==0.4.3
+sphinx>=3.0.0,!=3.5.0
+Jinja2<3.1

pyroengine/engine/engine.py

@@ -13,6 +13,7 @@
 from datetime import datetime, timedelta
 from collections import deque
 from typing import Optional, Dict
+import numpy as np
 
 from pyroclient import client
 from .predictor import PyronearPredictor
@@ -36,16 +37,18 @@ class PyronearEngine:
             the number of consecutive negative detections before stopping the alert
         cache_backup_period: number of minutes between each cache backup to disk
         frame_size: Resize frame to frame_size before sending it to the api in order to save bandwidth
+        model_weights: Path / url model yolov5 model weights
 
     Examples:
         >>> client_creds ={}
         >>> client_creds['pi_zero_id_1']={'login':'log1', 'password':'pwd1'}
         >>> client_creds['pi_zero_id_2']={'login':'log2', 'password':'pwd2'}
         >>> pyroEngine = PyronearEngine(0.6, 'https://api.pyronear.org', client_creds, 50)
     """
+
     def __init__(
         self,
-        detection_thresh: float = 0.5,
+        detection_thresh: float = 0.25,
         api_url: Optional[str] = None,
         client_creds: Optional[Dict[str, str]] = None,
         frame_saving_period: Optional[int] = None,
@@ -55,10 +58,12 @@ def __init__(
         alert_relaxation: int = 3,
         cache_backup_period: int = 60,
         frame_size: tuple = None,
+        model_weights: str = None,
     ) -> None:
         """Init engine"""
         # Engine Setup
-        self.pyronearPredictor = PyronearPredictor()
+
+        self.pyronearPredictor = PyronearPredictor(model_weights, detection_thresh)
         self.detection_thresh = detection_thresh
         self.frame_saving_period = frame_saving_period
         self.alert_relaxation = alert_relaxation
@@ -80,30 +85,42 @@ def __init__(
                 self.frames_counter[pi_zero_id] = 0
                 self.ongoing_alert[pi_zero_id] = False
         else:
-            self.consec_dets['-1'] = 0
-            self.ongoing_alert['-1'] = 0
+            self.consec_dets["-1"] = 0
+            self.ongoing_alert["-1"] = 0
 
         if self.api_url is not None:
             # Instantiate clients for each camera
             self.api_client = {}
             for _id, vals in client_creds.items():
-                self.api_client[_id] = client.Client(self.api_url, vals['login'], vals['password'])
+                self.api_client[_id] = client.Client(
+                    self.api_url, vals["login"], vals["password"]
+                )
 
         # Restore pending alerts cache
         self.pending_alerts = deque([], cache_size)
-        self._backup_folder = Path("data/")  # with Docker, the path has to be a bind volume
+        self._backup_folder = Path(
+            "data/"
+        )  # with Docker, the path has to be a bind volume
         self.load_cache_from_disk()
         self.cache_backup_period = cache_backup_period
         self.last_cache_dump = datetime.utcnow()
 
     def predict(self, frame: Image.Image, pi_zero_id: Optional[int] = None) -> float:
-        """ run prediction on comming frame"""
-        prob = self.pyronearPredictor.predict(frame.convert('RGB'))  # run prediction
+        """run prediction on comming frame"""
+        pred = self.pyronearPredictor.predict(frame.convert("RGB"))  # run prediction
+
+        if len(pred) > 0:
+            prob = np.max(pred[:, 4])
+        else:
+            prob = 0
+
         if pi_zero_id is None:
             logging.info(f"Wildfire detection score ({prob:.2%})")
         else:
             self.heartbeat(pi_zero_id)
-            logging.info(f"Wildfire detection score ({prob:.2%}), on device {pi_zero_id}")
+            logging.info(
+                f"Wildfire detection score ({prob:.2%}), on device {pi_zero_id}"
+            )
 
         # Reduce image size to save bandwidth
         if isinstance(self.frame_size, tuple):
@@ -112,10 +129,13 @@ def predict(self, frame: Image.Image, pi_zero_id: Optional[int] = None) -> float
         # Alert
         if prob > self.detection_thresh:
             if pi_zero_id is None:
-                pi_zero_id = '-1'  # add default key value
+                pi_zero_id = "-1"  # add default key value
 
             # Don't increment beyond relaxation
-            if not self.ongoing_alert[pi_zero_id] and self.consec_dets[pi_zero_id] < self.alert_relaxation:
+            if (
+                not self.ongoing_alert[pi_zero_id]
+                and self.consec_dets[pi_zero_id] < self.alert_relaxation
+            ):
                 self.consec_dets[pi_zero_id] += 1
 
             if self.consec_dets[pi_zero_id] == self.alert_relaxation:
@@ -144,15 +164,21 @@ def predict(self, frame: Image.Image, pi_zero_id: Optional[int] = None) -> float
             self.last_cache_dump = ts
 
         # save frame
-        if isinstance(self.api_url, str) and isinstance(self.frame_saving_period, int) and isinstance(pi_zero_id, int):
+        if (
+            isinstance(self.api_url, str)
+            and isinstance(self.frame_saving_period, int)
+            and isinstance(pi_zero_id, int)
+        ):
             self.frames_counter[pi_zero_id] += 1
             if self.frames_counter[pi_zero_id] == self.frame_saving_period:
                 # Reset frame counter
                 self.frames_counter[pi_zero_id] = 0
                 # Send frame to the api
-                frame.save(self.stream, format='JPEG')
+                frame.save(self.stream, format="JPEG")
                 self.save_frame(pi_zero_id)
-                self.stream.seek(0)  # "Rewind" the stream to the beginning so we can read its content
+                self.stream.seek(
+                    0
+                )  # "Rewind" the stream to the beginning so we can read its content
 
         return prob
 
@@ -162,16 +188,22 @@ def send_alert(self, pi_zero_id: int) -> None:
         # Create a media
         media_id = self.api_client[pi_zero_id].create_media_from_device().json()["id"]
         # Create an alert linked to the media and the event
-        self.api_client[pi_zero_id].send_alert_from_device(lat=self.latitude, lon=self.longitude, media_id=media_id)
-        self.api_client[pi_zero_id].upload_media(media_id=media_id, media_data=self.stream.getvalue())
+        self.api_client[pi_zero_id].send_alert_from_device(
+            lat=self.latitude, lon=self.longitude, media_id=media_id
+        )
+        self.api_client[pi_zero_id].upload_media(
+            media_id=media_id, media_data=self.stream.getvalue()
+        )
 
     def upload_frame(self, pi_zero_id: int) -> None:
         """Save frame"""
         logging.info("Uploading media...")
         # Create a media
         media_id = self.api_client[pi_zero_id].create_media_from_device().json()["id"]
         # Send media
-        self.api_client[pi_zero_id].upload_media(media_id=media_id, media_data=self.stream.getvalue())
+        self.api_client[pi_zero_id].upload_media(
+            media_id=media_id, media_data=self.stream.getvalue()
+        )
 
     def heartbeat(self, pi_zero_id: int) -> None:
         """Updates last ping of device"""
@@ -190,56 +222,68 @@ def upload_pending_alerts(self) -> None:
             frame_info = self.pending_alerts[0]
 
             try:
-                frame_info['frame'].save(self.stream, format='JPEG')
+                frame_info["frame"].save(self.stream, format="JPEG")
                 # Send alert to the api
-                self.send_alert(frame_info['pi_zero_id'])
+                self.send_alert(frame_info["pi_zero_id"])
                 # No need to upload it anymore
                 self.pending_alerts.popleft()
                 logging.info(f"Alert sent by device {frame_info['pi_zero_id']}")
             except ConnectionError:
-                logging.warning(f"Unable to upload cache for device {frame_info['pi_zero_id']}")
-                self.stream.seek(0)  # "Rewind" the stream to the beginning so we can read its content
+                logging.warning(
+                    f"Unable to upload cache for device {frame_info['pi_zero_id']}"
+                )
+                self.stream.seek(
+                    0
+                )  # "Rewind" the stream to the beginning so we can read its content
                 break
 
     def save_cache_to_disk(self) -> None:
 
         # Remove previous dump
-        json_path = self._backup_folder.joinpath('pending_alerts.json')
+        json_path = self._backup_folder.joinpath("pending_alerts.json")
         if json_path.is_file():
-            with open(json_path, 'rb') as f:
+            with open(json_path, "rb") as f:
                 data = json.load(f)
 
             for entry in data:
-                os.remove(entry['frame_path'])
+                os.remove(entry["frame_path"])
             os.remove(json_path)
 
         data = []
         for idx, info in enumerate(self.pending_alerts):
             # Save frame to disk
-            info['frame'].save(self._backup_folder.joinpath(f"pending_frame{idx}.jpg"))
+            info["frame"].save(self._backup_folder.joinpath(f"pending_frame{idx}.jpg"))
 
             # Save path in JSON
-            data.append({
-                "frame_path": str(self._backup_folder.joinpath(f"pending_frame{idx}.jpg")),
-                "pi_zero_id": info["pi_zero_id"],
-                "ts": info['ts']
-            })
+            data.append(
+                {
+                    "frame_path": str(
+                        self._backup_folder.joinpath(f"pending_frame{idx}.jpg")
+                    ),
+                    "pi_zero_id": info["pi_zero_id"],
+                    "ts": info["ts"],
+                }
+            )
 
         # JSON dump
         if len(data) > 0:
-            with open(json_path, 'w') as f:
+            with open(json_path, "w") as f:
                 json.dump(data, f)
 
     def load_cache_from_disk(self) -> None:
         # Read json
-        json_path = self._backup_folder.joinpath('pending_alerts.json')
+        json_path = self._backup_folder.joinpath("pending_alerts.json")
         if json_path.is_file():
-            with open(json_path, 'rb') as f:
+            with open(json_path, "rb") as f:
                 data = json.load(f)
 
             for entry in data:
                 # Open image
-                frame = Image.open(entry['frame_path'], mode='r')
+                frame = Image.open(entry["frame_path"], mode="r")
                 self.pending_alerts.append(
-                    {"frame": frame, "pi_zero_id": entry['pi_zero_id'], "ts": entry['ts']}
+                    {
+                        "frame": frame,
+                        "pi_zero_id": entry["pi_zero_id"],
+                        "ts": entry["ts"],
+                    }
                 )