6
0
Dimitri Korsch 3 жил өмнө
parent
commit
501f261d5d

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

@@ -22,6 +22,7 @@ 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)

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

@@ -1,3 +1,4 @@
+import chainer
 import numpy as np
 import typing as T
 
@@ -44,13 +45,15 @@ 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))
+        return str(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]