Add VideoDetector

environment.yml

@@ -5,7 +5,7 @@ channels:
   - conda-forge
   - defaults
 dependencies:
-  - python
+  - python <3.9
   - pip
   - pip:
     - -e .
@@ -19,16 +19,24 @@ dependencies:
   - wheel
 
   # Install Requirements (setup.py:requirements)
+  - av >=8.0.2,<9
   - click >=7.0
   - click-plugins
   - darknet
   - entrypoints
   - fsspec <=0.7.5
   - numpy
   - pillow
+  - tqdm
 
   # Zoo Optional Requirements
   - intake
+  - darknet-yolov4
+
+  # YouTube Optional Requirements
+  - pafy
+  - youtube-dl
+  - webvtt-py
 
   # Test Requirements (setup.py:test_requirements)
   - pytest >=3

examples/video.ipynb

@@ -0,0 +1,180 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "outputs": [],
+   "source": [
+    "from itertools import tee\n",
+    "import av\n",
+    "\n",
+    "import fsspec\n",
+    "\n",
+    "from darknet.py.detector import StreamDetector\n",
+    "from darknet.py.util import image_draw_detections as draw_detections\n"
+   ],
+   "metadata": {
+    "collapsed": false,
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   }
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "outputs": [],
+   "source": [
+    "darknet_gh_url = \"github://AlexeyAB:darknet@master\"\n",
+    "# Load the Coco labels/metadata\n",
+    "with fsspec.open(f\"{darknet_gh_url}/data/coco.names\", mode=\"rt\") as f:\n",
+    "    labels = [line.rstrip() for line in f.readlines()]\n",
+    "\n",
+    "d = StreamDetector(labels=labels,\n",
+    "                   config_url=f\"{darknet_gh_url}/cfg/yolov4.cfg\",\n",
+    "                   weights_url=\"https://github.com/AlexeyAB/darknet/releases/download/darknet_yolo_v3_optimal/yolov4.weights\",\n",
+    "                   batch_size=10)"
+   ],
+   "metadata": {
+    "collapsed": false,
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   }
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "outputs": [],
+   "source": [
+    "icntnr = av.open(\"ffmpeg/hell0.ts\")\n",
+    "istrm = icntnr.streams.video[0]\n",
+    "istrm.thread_type = \"AUTO\"\n",
+    "codec_context = icntnr.streams.video[0].codec_context\n",
+    "coded_height, coded_width = codec_context.coded_height, codec_context.coded_width\n"
+   ],
+   "metadata": {
+    "collapsed": false,
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   }
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 18,
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "CPU times: user 21 s, sys: 1.08 s, total: 22.1 s\n",
+      "Wall time: 8.25 s\n"
+     ]
+    }
+   ],
+   "source": [
+    "%%time\n",
+    "## We can't detect and trasncode in realtime.\n",
+    "## You should find a different solution if that is what you're trying to do.\n",
+    "icntnr.seek(0)\n",
+    "frames, detections = tee(icntnr.decode(video=0))\n",
+    "detections = d.detect(detections)\n",
+    "\n",
+    "ocntnr = av.open(\"ffmpeg/hell0-ann.ts\", mode=\"w\")\n",
+    "ostrm = ocntnr.add_stream(istrm.codec_context.codec.name, rate=istrm.codec_context.framerate)\n",
+    "ostrm.width = coded_width\n",
+    "ostrm.height = coded_height\n",
+    "\n",
+    "for frame, dets in zip(frames, detections):\n",
+    "    #print(f\"time={frame.time}, dts={frame.dts}, pts={frame.pts}: {dets}\")\n",
+    "    img = draw_detections(frame.to_image(), dets)\n",
+    "    oframe = av.VideoFrame.from_image(img)\n",
+    "    for packet in ostrm.encode(oframe):\n",
+    "        ocntnr.mux(packet)\n",
+    "\n",
+    "for packet in ostrm.encode():\n",
+    "    ocntnr.mux(packet)\n",
+    "\n",
+    "ocntnr.close()"
+   ],
+   "metadata": {
+    "collapsed": false,
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   }
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 19,
+   "outputs": [],
+   "source": [
+    "icntnr.close()"
+   ],
+   "metadata": {
+    "collapsed": false,
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   }
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 9,
+   "outputs": [
+    {
+     "data": {
+      "text/plain": "Fraction(30, 1)"
+     },
+     "execution_count": 9,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "icntnr.start_time\n",
+    "icntnr.streams.video[0]."
+   ],
+   "metadata": {
+    "collapsed": false,
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   }
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "outputs": [],
+   "source": [],
+   "metadata": {
+    "collapsed": false,
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   }
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 2
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython2",
+   "version": "2.7.6"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 0
+}

examples/video.py

@@ -0,0 +1,99 @@
+from copy import deepcopy
+from itertools import tee
+import av
+import click
+import fsspec
+import json
+import m3u8
+import pafy
+from tqdm import tqdm
+import webvtt
+
+from darknet.py.detector import StreamDetector
+
+
+def dets_to_json(dets):
+    return json.dumps(
+        {
+            "detections": [
+                {"label": det[0], "confidence": int(det[1] * 100), "bbox": [int(c) for c in det[2]]}
+                for det in dets
+            ]
+        }
+    )
+
+
+@click.command()
+@click.argument("playlist-uri", default="https://www.youtube.com/watch?v=iLQNbwMWbGM")
+@click.option("--batch-size", "-b", default=4, help="inference batch size (use power of 2)")
+@click.option("--labels-uri", "-l", default="github://AlexeyAB:darknet@master/data/coco.names")
+@click.option("--config-uri", "-c", default="github://AlexeyAB:darknet@master/cfg/yolov4.cfg")
+@click.option("--weights-uri", "-w", default="https://github.com/AlexeyAB/darknet/releases/download/darknet_yolo_v3_optimal/yolov4.weights")
+def main(batch_size, playlist_uri, labels_uri, config_uri, weights_uri):
+    """
+    Creates WebVTT metadata for objects in frames.
+
+    """
+    darknet_gh_url = "github://AlexeyAB:darknet@master"
+    # Load the Coco labels/metadata
+    with fsspec.open(labels_uri, mode="rt") as f:
+        labels = [line.rstrip() for line in f.readlines()]
+
+    model = ".".join(weights_uri.split("/")[-1].split(".")[:-1])
+    d = StreamDetector(
+        labels=labels,
+        config_url=config_uri,
+        weights_url=weights_uri,
+        batch_size=batch_size,
+    )
+
+    if playlist_uri.startswith("https://www.youtube.com/watch?v="):
+        playlist_uri = pafy.new(playlist_uri).getbest().url
+
+    playlist = m3u8.load(playlist_uri)
+
+    seg_playlist = deepcopy(playlist)
+    seg_playlist.segments.clear()
+    seg_playlist.is_endlist = False
+
+    vtt_playlist = deepcopy(playlist)
+    vtt_playlist.segments.clear()
+    vtt_playlist.is_endlist = False
+
+    for idx, segment in enumerate(tqdm(playlist.segments, leave=True, desc=playlist_uri)):
+        seg_playlist.add_segment(segment)
+        vtt_segment = deepcopy(segment)
+        vtt_segment.uri = f"{model}-{idx}.vtt"
+
+        with av.open(segment.absolute_uri) as icntnr:
+            istrm = icntnr.streams.video[0]
+            istrm.thread_type = "AUTO"
+
+            frames, detections = tee(icntnr.decode(video=0))
+            detections = d.detect(detections)
+
+            vtt = webvtt.WebVTT()
+            for frame, dets in tqdm(zip(frames, detections), desc=vtt_segment.uri, leave=True):
+                if len(dets) == 0:
+                    continue
+                caption = webvtt.Caption()
+                caption.text = dets_to_json(dets)
+                caption._start = frame.time
+                caption._end = caption._start + 1 / istrm.framerate
+                vtt.captions.append(caption)
+
+            vtt.save(vtt_segment.uri)
+            vtt_playlist.add_segment(vtt_segment)
+
+            vtt_playlist.dump(f"{model}.m3u8")
+            seg_playlist.dump(f"{model}-sync.m3u8")
+
+    vtt_playlist.is_endlist = True
+    vtt_playlist.dump(f"{model}.m3u8")
+
+    seg_playlist.is_endlist = True
+    seg_playlist.dump(f"{model}-sync.m3u8")
+
+
+if __name__ == "__main__":
+    main()

setup.py

@@ -11,6 +11,7 @@
 # The requirements section should be kept in sync with the environment.yml file
 requirements = [
     # fmt: off
+    "av >=8.0.2,<9",
     "click>=7.0",
     "click-plugins",
     "entrypoints",

src/darknet/__init__.py

@@ -1,2 +1 @@
-# This is here for Python 2.x compatibility
-__path__ = __import__("pkgutil").extend_path(__path__, __name__)
+__import__("pkg_resources").declare_namespace(__name__)

src/darknet/py/__init__.py

@@ -1,6 +1,6 @@
 """Top-level package for DarkNet OpenSource Neural Networks in Python."""
 from ._version import version as __version__  # noqa: F401
 from .classifier import Classifier, ImageClassifier
-from .detector import ImageDetector
+from .detector import ImageDetector, StreamDetector
 
-__all__ = ["Classifier", "ImageClassifier", "ImageDetector"]
+__all__ = ["Classifier", "ImageClassifier", "ImageDetector", "StreamDetector"]

src/darknet/py/cli.py

@@ -14,6 +14,3 @@ def py(args=None):
     # fmt: on
     return 0
 
-
-if __name__ == "__main__":
-    sys.exit(py)  # pragma: no cover