hysts HF Staff commited on
Commit
a9e8a09
·
1 Parent(s): d38fcf4

Add Additional tag input

Browse files
Files changed (1) hide show
  1. app.py +9 -2
app.py CHANGED
@@ -73,6 +73,7 @@ def run(
73
  hair_color_tags: list[str],
74
  hair_style_tags: list[str],
75
  image_color_tags: list[str],
 
76
  score_threshold: float,
77
  start_index: int,
78
  nrows: int,
@@ -84,11 +85,16 @@ def run(
84
  tarball_path: pathlib.Path,
85
  deepdanbooru_tag_dict: dict[str, int],
86
  deepdanbooru_predictions: np.ndarray,
87
- ) -> np.ndarray:
88
  hair_color_tags = [f'{color}_hair' for color in hair_color_tags]
 
89
 
90
  tags = general_tags + hair_color_tags + hair_style_tags + image_color_tags
91
- tag_indices = [deepdanbooru_tag_dict[tag] for tag in tags]
 
 
 
 
92
 
93
  conditions = deepdanbooru_predictions[:, tag_indices] > score_threshold
94
  image_indices = np.arange(len(deepdanbooru_predictions))
@@ -190,6 +196,7 @@ def main():
190
  'monochrome',
191
  ],
192
  label='Image Color'),
 
193
  gr.inputs.Slider(0,
194
  1,
195
  step=0.1,
 
73
  hair_color_tags: list[str],
74
  hair_style_tags: list[str],
75
  image_color_tags: list[str],
76
+ additional_tags: str,
77
  score_threshold: float,
78
  start_index: int,
79
  nrows: int,
 
85
  tarball_path: pathlib.Path,
86
  deepdanbooru_tag_dict: dict[str, int],
87
  deepdanbooru_predictions: np.ndarray,
88
+ ) -> tuple[int, np.ndarray, np.ndarray]:
89
  hair_color_tags = [f'{color}_hair' for color in hair_color_tags]
90
+ additional_tags = additional_tags.split(',')
91
 
92
  tags = general_tags + hair_color_tags + hair_style_tags + image_color_tags
93
+ tags += additional_tags
94
+ tag_indices = [
95
+ deepdanbooru_tag_dict[tag] for tag in tags
96
+ if tag in deepdanbooru_tag_dict
97
+ ]
98
 
99
  conditions = deepdanbooru_predictions[:, tag_indices] > score_threshold
100
  image_indices = np.arange(len(deepdanbooru_predictions))
 
196
  'monochrome',
197
  ],
198
  label='Image Color'),
199
+ gr.inputs.Textbox(label='Additional Tags'),
200
  gr.inputs.Slider(0,
201
  1,
202
  step=0.1,