|
@@ -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:
|