Use fsspec urls for opening the darknet network

Author: sodre

environment.yml

@@ -31,6 +31,7 @@ dependencies:
   # Test Requirements (setup.py:test_requirements)
   - pytest >=3
   - pytest-cov
+  - pytest-mock
 
   # Documentation Requirements (setup.py:doc_requirements)
   - sphinx

examples/coco.ipynb

@@ -2,7 +2,7 @@
  "cells": [
   {
    "cell_type": "code",
-   "execution_count": 10,
+   "execution_count": 1,
    "metadata": {
     "collapsed": true
    },
@@ -18,27 +18,25 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 3,
+   "execution_count": 2,
    "metadata": {
     "pycharm": {
      "name": "#%%\n"
     }
    },
    "outputs": [],
    "source": [
-    "github_opts = dict(org=\"pjreddie\", repo=\"darknet\", sha=\"master\")\n",
-    "darknet_gh_url = \"github://{org}:{repo}@master\".format(**github_opts)\n",
-    "filecache_opts = dict(cache_storage=f\"{os.environ['HOME']}/.cache/darknet.py\")"
+    "darknet_gh_url = \"github://pjreddie:darknet@master\""
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 4,
+   "execution_count": 3,
    "outputs": [],
    "source": [
     "# Load the ImageNet 1k labels/metadata\n",
-    "with fsspec.open(f\"{darknet_gh_url}/data/imagenet.shortnames.list\") as f:\n",
-    "    labels = [line.decode().rstrip() for line in f.readlines()[:1000]]"
+    "with fsspec.open(f\"{darknet_gh_url}/data/imagenet.shortnames.list\", mode=\"rt\") as f:\n",
+    "    labels = [line.rstrip() for line in f.readlines()[:1000]]"
    ],
    "metadata": {
     "collapsed": false,
@@ -49,29 +47,22 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 6,
+   "execution_count": 5,
    "metadata": {
     "pycharm": {
      "name": "#%%\n"
     }
    },
    "outputs": [],
    "source": [
-    "config_of = fsspec.open(f\"filecache::github://cfg/darknet53_448.cfg\",\n",
-    "                     filecache=filecache_opts, github=github_opts)\n",
-    "weights_of = fsspec.open(\"filecache::https://pjreddie.com/media/files/darknet53_448.weights\",\n",
-    "                      filecache=filecache_opts)\n",
-    "with config_of as config:\n",
-    "    with weights_of as weights:\n",
-    "        # Load the Classifier\n",
-    "        n = ImageClassifier(labels=labels,\n",
-    "                            config_file=config.name,\n",
-    "                            weights_file=weights.name)"
+    "n = ImageClassifier(labels=labels,\n",
+    "                    config_url=f\"{darknet_gh_url}/cfg/darknet53_448.cfg\",\n",
+    "                    weights_url=\"https://pjreddie.com/media/files/darknet53_448.weights\")"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 7,
+   "execution_count": 6,
    "metadata": {
     "pycharm": {
      "name": "#%%\n"
@@ -82,7 +73,7 @@
      "data": {
       "text/plain": "[('malamute', 0.98354006),\n ('Eskimo dog', 0.0042837244),\n ('Siberian husky', 0.0031863458),\n ('Tibetan mastiff', 0.0030448402),\n ('Great Pyrenees', 0.0022190544)]"
      },
-     "execution_count": 7,
+     "execution_count": 6,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -94,13 +85,13 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 12,
+   "execution_count": 7,
    "outputs": [
     {
      "data": {
       "text/plain": "[('malamute', 0.98354006),\n ('Eskimo dog', 0.0042837244),\n ('Siberian husky', 0.0031863458),\n ('Tibetan mastiff', 0.0030448402),\n ('Great Pyrenees', 0.0022190544)]"
      },
-     "execution_count": 12,
+     "execution_count": 7,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -119,7 +110,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 13,
+   "execution_count": 8,
    "metadata": {
     "pycharm": {
      "name": "#%%\n"
@@ -130,7 +121,7 @@
      "data": {
       "text/plain": "[('malamute', 0.98354006),\n ('Eskimo dog', 0.0042837244),\n ('Siberian husky', 0.0031863458),\n ('Tibetan mastiff', 0.0030448402),\n ('Great Pyrenees', 0.0022190544)]"
      },
-     "execution_count": 13,
+     "execution_count": 8,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -150,7 +141,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 14,
+   "execution_count": 9,
    "metadata": {
     "pycharm": {
      "name": "#%%\n"
@@ -161,7 +152,7 @@
      "data": {
       "text/plain": "[('bald eagle', 0.55666465),\n ('vulture', 0.21876547),\n ('kite', 0.18937683),\n ('ruddy turnstone', 0.004589723),\n ('ruffed grouse', 0.0032499917)]"
      },
-     "execution_count": 14,
+     "execution_count": 9,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -172,7 +163,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 15,
+   "execution_count": 10,
    "metadata": {
     "pycharm": {
      "name": "#%%\n"
@@ -183,7 +174,7 @@
      "data": {
       "text/plain": "[('electric guitar', 0.98759043),\n ('acoustic guitar', 0.009553942),\n ('banjo', 0.0011607071),\n ('pick', 0.0007309786),\n ('stage', 0.00058993115)]"
      },
-     "execution_count": 15,
+     "execution_count": 10,
      "metadata": {},
      "output_type": "execute_result"
     }

examples/imagenet1k-intake.ipynb

@@ -35,6 +35,7 @@
     # fmt: off
     "pytest>=3",
     "pytest-cov",
+    "pytest-mock",
     # fmt: on
 ]
 

examples/imagenet1k.ipynb

@@ -10,8 +10,8 @@ class ClassifierBase(ABC):
     network: Network
     labels: list
 
-    def __init__(self, labels, config_file, weights_file, **kwargs):
-        self.network = Network(config_file, weights_file, **kwargs)
+    def __init__(self, labels, config_url, weights_url, **kwargs):
+        self.network = Network.open(config_url, weights_url, **kwargs)
 
         self.labels = range(self.network.output_size()) if labels is None else labels
         if len(self.labels) != self.network.output_size():

setup.py

@@ -8,8 +8,8 @@ class ImageDetector(object):
 
     _last_image_size = None
 
-    def __init__(self, labels, config_file, weights_file, **kwargs):
-        self.network = Network(config_file, weights_file, **kwargs)
+    def __init__(self, labels, config_url, weights_url, **kwargs):
+        self.network = Network.open(config_url, weights_url, **kwargs)
         self.labels = labels
 
     def detect(self, image, **kwargs):

src/darknet/py/classifier.py

@@ -6,6 +6,7 @@ cimport numpy as np
 cimport libdarknet as dn
 
 from libc.stdlib cimport free
+from .util import fsspec_cache_open
 
 np.import_array()
 
@@ -24,6 +25,13 @@ cdef class Metadata:
 cdef class Network:
     cdef dn.network* _c_network
 
+    @staticmethod
+    def open(config_url, weights_url):
+        with fsspec_cache_open(config_url, mode="rt") as config:
+            with fsspec_cache_open(weights_url, mode="rb") as weights:
+                return Network(config.name, weights.name)
+
+
     def __cinit__(self, config_file, weights_file):
         clear = 1
         self._c_network = dn.load_network(config_file.encode(), weights_file.encode(), clear)
@@ -96,7 +104,7 @@ cdef class Network:
             elif nms_type == "sort":
                 dn.do_nms_sort(detections, num_dets, detections[0].classes, nms_threshold)
             else:
-                raise ValueError(f"non-maximum-supression type {nms_type} is not one of {['obj', 'sort']}")
+                raise ValueError(f"non-maximum-suppression type {nms_type} is not one of {['obj', 'sort']}")
 
         rv = [
             (j, detections[i].prob[j],

src/darknet/py/detector.py

@@ -1,8 +1,55 @@
+import os
+import fsspec
+
 import numpy as np
 from PIL import Image
 from PIL import ImageDraw
 
-import fsspec
+
+def fsspec_cache_open(urlpath: str,
+                      mode="rb",
+                      compression=None,
+                      encoding="utf8",
+                      errors=None,
+                      protocol=None,
+                      newline=None,
+                      **kwargs) -> fsspec.core.OpenFile:
+    chain = urlpath.split("::")
+
+    if chain[0].startswith("github"):
+        chain[0], kwargs = fsspec_split_github_url(chain[0], kwargs)
+
+    # Because darknet is written in C, we need real file names for it to open
+    if chain[0] not in {"filecache", "simplecache"}:
+        first_scheme = chain[0].split("://")[0]
+        urlpath = f"filecache::{urlpath}"
+        filecache = dict(cache_storage=f"{os.environ['HOME']}/.cache/darknet.py")
+        kwargs = {
+            "filecache": filecache,
+            first_scheme: kwargs
+        }
+
+    return fsspec.open(urlpath, mode, compression, encoding, errors, protocol, newline, **kwargs)
+
+
+def fsspec_split_github_url(github_url: str, kwargs: dict) -> (str, dict):
+    # TODO: Remove this once fsspec > 0.7.5
+    from urllib.parse import urlparse
+
+    rv = github_url
+    github_url = urlparse(github_url)
+    keys = {"org", "repo", "sha"}
+    # If that metadata is not passed as kwargs, we need to extract it
+    # netloc = "{org}:{repo}@{sha}"
+    kwargs = kwargs or dict()
+    if (keys & kwargs.keys()) != keys:
+        org, repo, sha = github_url.username, github_url.password, github_url.hostname
+        if org is None or repo is None or sha is None:
+            raise ValueError(f"The github url {github_url} does not match `github://<org>:<repo>@<sha>/path`")
+        kwargs.update(dict(org=org, repo=repo, sha=sha))
+        rv = github_url.geturl().replace(f"{github_url.netloc}/", "")
+
+    return rv, kwargs
 
 
 def image_to_3darray(image, target_shape):

src/darknet/py/network.pyx

@@ -0,0 +1,20 @@
+import fsspec
+import darknet.py.util as darknet_util
+
+def test_fsspec_split_github_url():
+    url, kwargs = darknet_util.fsspec_split_github_url("github://org:repo@sha/path/to/file", None)
+    assert url == "github://path/to/file"
+    assert kwargs["org"] == "org"
+    assert kwargs["repo"] == "repo"
+    assert kwargs["sha"] == "sha"
+
+
+def test_fsspec_cache_open(mocker):
+    split_spy = mocker.spy(darknet_util, "fsspec_split_github_url")
+    fsspec_open_mock = mocker.patch.object(fsspec, "open", autospec=True)
+
+    of = darknet_util.fsspec_cache_open("github://org:repo@sha/path/to/file")
+    assert fsspec_open_mock.return_value == of
+
+    fsspec_open_mock.assert_called_once()
+    split_spy.assert_called_once()