Irina Tolstykh commited on
Commit
f4e3f66
1 Parent(s): 5f16544

add inference paramters

Browse files
Files changed (1) hide show
  1. app.py +33 -3
app.py CHANGED
@@ -50,9 +50,31 @@ def load_models():
50
 
51
  return predictor
52
 
53
- def detect(image: np.ndarray, predictor: Predictor) -> np.ndarray:
 
 
 
 
 
 
54
  # input is rgb image, output must be rgb too
55
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
56
  image = image[:, :, ::-1] # RGB -> BGR
57
  detected_objects, out_im = predictor.recognize(image)
58
  return out_im[:, :, ::-1] # BGR -> RGB
@@ -61,13 +83,21 @@ def detect(image: np.ndarray, predictor: Predictor) -> np.ndarray:
61
  predictor = load_models()
62
 
63
  image_dir = pathlib.Path('images')
64
- examples = [[path.as_posix()] for path in sorted(image_dir.glob('*.jpg'))]
65
 
66
  func = functools.partial(detect, predictor=predictor)
67
 
68
  gr.Interface(
69
  fn=func,
70
- inputs=gr.Image(label='Input', type='numpy'),
 
 
 
 
 
 
 
 
71
  outputs=gr.Image(label='Output', type='numpy'),
72
  examples=examples,
73
  examples_per_page=30,
 
50
 
51
  return predictor
52
 
53
+ def detect(
54
+ image: np.ndarray,
55
+ score_threshold: float,
56
+ iou_threshold: float,
57
+ mode: str,
58
+ predictor: Predictor
59
+ ) -> np.ndarray:
60
  # input is rgb image, output must be rgb too
61
 
62
+ predictor.detector.detector_kwargs['conf'] = score_threshold
63
+ predictor.detector.detector_kwargs['iou'] = iou_threshold
64
+
65
+ if mode == "Use persons and faces":
66
+ use_persons = True
67
+ disable_faces = False
68
+ elif mode == "Use persons only":
69
+ use_persons = True
70
+ disable_faces = True
71
+ elif mode == "Use faces only":
72
+ use_persons = False
73
+ disable_faces = False
74
+
75
+ predictor.age_gender_model.meta.use_persons = use_persons
76
+ predictor.age_gender_model.meta.disable_faces = disable_faces
77
+
78
  image = image[:, :, ::-1] # RGB -> BGR
79
  detected_objects, out_im = predictor.recognize(image)
80
  return out_im[:, :, ::-1] # BGR -> RGB
 
83
  predictor = load_models()
84
 
85
  image_dir = pathlib.Path('images')
86
+ examples = [[path.as_posix(), 0.4, 0.7, "Use persons and faces"] for path in sorted(image_dir.glob('*.jpg'))]
87
 
88
  func = functools.partial(detect, predictor=predictor)
89
 
90
  gr.Interface(
91
  fn=func,
92
+ inputs=[
93
+ gr.Image(label='Input', type='numpy'),
94
+ gr.Slider(0, 1, value=0.4, step=0.05, label='Score Threshold'),
95
+ gr.Slider(0, 1, value=0.7, step=0.05, label='Iou Threshold'),
96
+ gr.Radio(["Use persons and faces", "Use persons only", "Use faces only"],
97
+ value="Use persons and faces",
98
+ label="Inference mode",
99
+ info="What to use for gender and age recognition"),
100
+ ],
101
  outputs=gr.Image(label='Output', type='numpy'),
102
  examples=examples,
103
  examples_per_page=30,