소스 검색

added first code for the moth classifier (not tested yet)

Dimitri Korsch 4 년 전
부모
커밋
b33401feaa
3개의 변경된 파일44개의 추가작업 그리고 0개의 파일을 삭제
  1. 5 0
      models/moth_scanner/configuration.json
  2. 36 0
      models/moth_scanner/scanner/classifier.py
  3. 3 0
      requirements.txt

+ 5 - 0
models/moth_scanner/configuration.json

@@ -19,5 +19,10 @@
       "dilate_iterations": 3,
       "kernel_size": 5
     }
+  },
+  "classifier": {
+    "model_type": "cvmodelz.InceptionV3",
+    "weights": "classifier.npz",
+    "n_classes": 200,
   }
 }

+ 36 - 0
models/moth_scanner/scanner/classifier.py

@@ -0,0 +1,36 @@
+import numpy as np
+import typing as T
+
+from munch import munchify
+
+from cvmodelz.models import ModelFactory
+
+
+def
+
+class Classifier(object):
+
+    def __init__(self, configuration: T.Dict):
+        super().__init__()
+
+        config = munchify(configuration)
+
+        model_type = config.model_type
+        n_classes = config.n_classes
+        weights = config.weights
+
+        self.backbone = ModelFactory.new(model_type)
+        self.backbone.load_for_inference(weights,
+                                         n_classes=n_classes,
+                                         path="model/",
+                                         strict=True,
+                                        )
+
+    def __call__(self, im: np.ndarray):
+        assert im.ndim in (3, 4), \
+            "Classifier accepts only RGB images (3D input) or a batch of images (4D input)!"
+
+        if im.ndim == 3:
+            # expand first axis
+            im = im[None]
+

+ 3 - 0
requirements.txt

@@ -7,3 +7,6 @@ flask
 python-socketio
 munch
 scikit-image
+
+chainer~=7.8
+cvmodelz~=0.1