Irina Tolstykh
commited on
Commit
•
025ec32
1
Parent(s):
94e188b
fix channels
Browse files
app.py
CHANGED
@@ -26,12 +26,15 @@ def load_model():
|
|
26 |
return yolo
|
27 |
|
28 |
def detect(image: np.ndarray, detector: YOLO) -> np.ndarray:
|
|
|
|
|
29 |
detector_kwargs = {'conf': 0.5, 'iou': 0.5, 'half': False, 'verbose': False}
|
30 |
|
|
|
31 |
results: Results = detector.predict(image, **detector_kwargs)[0]
|
32 |
out_im = results.plot()
|
33 |
|
34 |
-
return out_im
|
35 |
|
36 |
|
37 |
detector = load_model()
|
|
|
26 |
return yolo
|
27 |
|
28 |
def detect(image: np.ndarray, detector: YOLO) -> np.ndarray:
|
29 |
+
# input is rgb image, output must be rgb too
|
30 |
+
|
31 |
detector_kwargs = {'conf': 0.5, 'iou': 0.5, 'half': False, 'verbose': False}
|
32 |
|
33 |
+
image = image[:, :, ::-1] # RGB -> BGR
|
34 |
results: Results = detector.predict(image, **detector_kwargs)[0]
|
35 |
out_im = results.plot()
|
36 |
|
37 |
+
return out_im[:, :, ::-1] # BGR -> RGB
|
38 |
|
39 |
|
40 |
detector = load_model()
|