Irina Tolstykh commited on
Commit
025ec32
1 Parent(s): 94e188b

fix channels

Browse files
Files changed (1) hide show
  1. app.py +4 -1
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()