瀏覽代碼

added flat model label reader

Dimitri Korsch 3 年之前
父節點
當前提交
7e8506c04f

+ 5 - 1
models/moth_scanner/scanner/__init__.py

@@ -22,15 +22,19 @@ class Scanner(Interface):
     def execute(self, storage: MediaStorage, file: MediaFile):
 
         im = self.read_image(file.path)
+        im = cv2.cvtColor(im, cv2.COLOR_BGR2RGB)
         bw_im = cv2.cvtColor(im, cv2.COLOR_RGB2GRAY)
 
         detections = self.detector(bw_im)
 
+        labels = {ml.reference: ml for ml in storage.labels()}
+
         for bbox, info in detections:
             if not info.selected:
                 continue
             x0, y0, x1, y1 = bbox
-            label = self.classifier(bbox.crop(im, enlarge=True))
+            cls_id = self.classifier(bbox.crop(im, enlarge=True))
+            label = labels.get(str(cls_id), cls_id)
             file.add_bounding_box(x0, y0, bbox.w, bbox.h, label=label)
 
     def read_image(self, path: str, mode: int = cv2.IMREAD_COLOR) -> np.ndarray:

+ 5 - 3
models/moth_scanner/scanner/classifier.py

@@ -1,3 +1,4 @@
+import chainer
 import numpy as np
 import typing as T
 
@@ -44,13 +45,14 @@ class Classifier(object):
 
         if im.ndim == 3:
             # expand first axis
-            # CxHxW -> 1xCxHxW
+            # HxWxC -> 1xHxWxC
             im = im[None]
 
+
         im = [self._transform(_im) for _im in im]
         x = self.backbone.xp.array(im)
-        pred = self.backbone(x)
+        with chainer.using_config("train", False), chainer.no_backprop_mode():
+            pred = self.backbone(x)
         pred.to_cpu()
 
         return int(np.argmax(pred.array, axis=1))
-

+ 4 - 3
models/moth_scanner/scanner/detector.py

@@ -43,11 +43,12 @@ class BBox(namedtuple("BBox", "x0 y0 x1 y1")):
 
         # enlarge to a square extent
         if enlarge:
-            h, w = int(self.h * H), int(self.h * W)
+            h, w = int(self.h * H), int(self.w * W)
             size = max(h, w)
             dw, dh = (size - w) / 2, (size - h) / 2
-            x0, x1 = max(int(x0 - dw), 0), int(x0 - dw + size)
-            y0, y1 = max(int(y0 - dh), 0), int(y0 - dh + size)
+            x0, y0 = max(int(x0 - dw), 0), max(int(y0 - dh), 0)
+            x1, y1 = int(x0 + size), int(y0 + size)
+
 
         if im.ndim == 2:
             return im[y0:y1, x0:x1]