Haiyu Wu commited on
Commit
28a6a8a
·
1 Parent(s): 7aa5b58
Files changed (1) hide show
  1. app.py +55 -16
app.py CHANGED
@@ -10,6 +10,23 @@ from sixdrepnet.model import SixDRepNet
10
  import pixel_generator.vec2face.model_vec2face as model_vec2face
11
  MAX_SEED = np.iinfo(np.int32).max
12
  import torch
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
13
 
14
 
15
  def sample_nearby_vectors(base_vector, epsilons=[0.3, 0.5, 0.7], percentages=[0.4, 0.4, 0.2]):
@@ -62,7 +79,7 @@ def initialize_models():
62
  return generator, id_model, pose_model, quality_model
63
 
64
 
65
- def image_generation(input_image, quality, use_target_pose, pose, dimension):
66
  generator, id_model, pose_model, quality_model = initialize_models()
67
 
68
  generated_images = []
@@ -77,12 +94,10 @@ def image_generation(input_image, quality, use_target_pose, pose, dimension):
77
  if not use_target_pose:
78
  features = []
79
  norm = np.linalg.norm(feature, 2, 1, True)
80
- for i in np.arange(0, 4.8, 0.8):
81
  updated_feature = feature
82
  updated_feature[0][dimension] = feature[0][dimension] + i
83
-
84
  updated_feature = updated_feature / np.linalg.norm(updated_feature, 2, 1, True) * norm
85
-
86
  features.append(updated_feature)
87
  features = torch.tensor(np.vstack(features)).float()
88
  if quality > 25:
@@ -90,7 +105,7 @@ def image_generation(input_image, quality, use_target_pose, pose, dimension):
90
  else:
91
  _, _, images, *_ = generator(features)
92
  else:
93
- features = torch.repeat_interleave(torch.tensor(feature), 6, dim=0)
94
  features = sample_nearby_vectors(features, [0.7], [1]).float()
95
  if quality > 25 or pose > 20:
96
  images, _ = generator.gen_image(features, quality_model, id_model, pose_model=pose_model,
@@ -99,12 +114,13 @@ def image_generation(input_image, quality, use_target_pose, pose, dimension):
99
  _, _, images, *_ = generator(features)
100
 
101
  images = ((images.permute(0, 2, 3, 1).detach().cpu().numpy() + 1) / 2 * 255).astype(np.uint8)
102
- for image in images:
103
  generated_images.append(Image.fromarray(image))
 
104
  return generated_images
105
 
106
 
107
- def process_input(image_input, num1, num2, num3, num4, random_seed, target_quality, use_target_pose, target_pose):
108
  # Ensure all dimension numbers are within [0, 512)
109
  num1, num2, num3, num4 = [max(0, min(int(n), 511)) for n in [num1, num2, num3, num4]]
110
 
@@ -118,13 +134,14 @@ def process_input(image_input, num1, num2, num3, num4, random_seed, target_quali
118
  input_data = Image.open(image_input)
119
  input_data = np.array(input_data.resize((112, 112)))
120
 
121
- generated_images = image_generation(input_data, target_quality, use_target_pose, target_pose, [num1, num2, num3, num4])
122
 
123
  return generated_images
124
 
 
125
  def select_image(value, images):
126
  # Convert the float value (0 to 4) to an integer index (0 to 9)
127
- index = int(value / 0.8)
128
  return images[index]
129
 
130
  def toggle_inputs(use_pose):
@@ -147,9 +164,9 @@ def main():
147
  <b>Official 🤗 Gradio demo</b> for <a href='https://github.com/HaiyuWu/vec2face' target='_blank'><b>Vec2Face: Scaling Face Dataset Generation with Loosely Constrained Vectors</b></a>.<br>
148
 
149
  How to use:<br>
150
- 1. Upload an image with a cropped face image or directly click <b>Submit</b> button, six images will be shown on the right.
151
  2. You can control the image quality, image pose, and modify the values in the target dimensions to change the output images.
152
- 3. The output results will shown six results of dimension modification or pose images.
153
  4. Since the demo is CPU-based, higher quality and larger pose need longer time to run.
154
  5. Enjoy! 😊
155
  """
@@ -167,9 +184,9 @@ def main():
167
 
168
  with gr.Row():
169
  num1 = gr.Number(label="Dimension 1", value=0, minimum=0, maximum=511, step=1)
170
- num2 = gr.Number(label="Dimension 2", value=0, minimum=0, maximum=511, step=1)
171
- num3 = gr.Number(label="Dimension 3", value=0, minimum=0, maximum=511, step=1)
172
- num4 = gr.Number(label="Dimension 4", value=0, minimum=0, maximum=511, step=1)
173
 
174
  random_seed = gr.Number(label="Random Seed", value=42, minimum=0, maximum=MAX_SEED, step=1)
175
  target_quality = gr.Slider(label="Minimum Quality", minimum=22, maximum=35, step=1, value=24)
@@ -191,9 +208,10 @@ def main():
191
 
192
  with gr.Column():
193
  gallery = gr.Image(label="Generated Image")
 
194
  incremental_value_slider = gr.Slider(
195
  label="Result of dimension modification or results of pose images",
196
- minimum=0, maximum=4, step=0.8, value=0
197
  )
198
  gr.Markdown("""
199
  - These values are added to the dimensions (before normalization), **please ignore it if pose editing is on**.
@@ -208,14 +226,35 @@ def main():
208
  generated_images = gr.State([])
209
 
210
  submit.click(
 
 
 
 
 
 
 
 
211
  fn=process_input,
212
  inputs=[image_file, num1, num2, num3, num4, random_seed, target_quality, use_target_pose, target_pose],
213
  outputs=[generated_images]
 
 
 
 
214
  ).then(
215
  fn=select_image,
216
  inputs=[incremental_value_slider, generated_images],
217
  outputs=[gallery]
218
  )
 
 
 
 
 
 
 
 
 
219
 
220
  incremental_value_slider.change(
221
  fn=select_image,
@@ -244,4 +283,4 @@ def main():
244
 
245
 
246
  if __name__ == "__main__":
247
- main()
 
10
  import pixel_generator.vec2face.model_vec2face as model_vec2face
11
  MAX_SEED = np.iinfo(np.int32).max
12
  import torch
13
+ from time import time
14
+
15
+
16
+ def clear_image():
17
+ return None
18
+
19
+
20
+ def clear_generation_time():
21
+ return ""
22
+
23
+
24
+ def generating():
25
+ return "Generating images..."
26
+
27
+
28
+ def done():
29
+ return "Done!"
30
 
31
 
32
  def sample_nearby_vectors(base_vector, epsilons=[0.3, 0.5, 0.7], percentages=[0.4, 0.4, 0.2]):
 
79
  return generator, id_model, pose_model, quality_model
80
 
81
 
82
+ def image_generation(input_image, quality, use_target_pose, pose, dimension, progress=gr.Progress()):
83
  generator, id_model, pose_model, quality_model = initialize_models()
84
 
85
  generated_images = []
 
94
  if not use_target_pose:
95
  features = []
96
  norm = np.linalg.norm(feature, 2, 1, True)
97
+ for i in progress.tqdm(np.arange(0, 4.8, 2), desc="Generating images"):
98
  updated_feature = feature
99
  updated_feature[0][dimension] = feature[0][dimension] + i
 
100
  updated_feature = updated_feature / np.linalg.norm(updated_feature, 2, 1, True) * norm
 
101
  features.append(updated_feature)
102
  features = torch.tensor(np.vstack(features)).float()
103
  if quality > 25:
 
105
  else:
106
  _, _, images, *_ = generator(features)
107
  else:
108
+ features = torch.repeat_interleave(torch.tensor(feature), 3, dim=0)
109
  features = sample_nearby_vectors(features, [0.7], [1]).float()
110
  if quality > 25 or pose > 20:
111
  images, _ = generator.gen_image(features, quality_model, id_model, pose_model=pose_model,
 
114
  _, _, images, *_ = generator(features)
115
 
116
  images = ((images.permute(0, 2, 3, 1).detach().cpu().numpy() + 1) / 2 * 255).astype(np.uint8)
117
+ for image in progress.tqdm(images, desc="Processing images"):
118
  generated_images.append(Image.fromarray(image))
119
+
120
  return generated_images
121
 
122
 
123
+ def process_input(image_input, num1, num2, num3, num4, random_seed, target_quality, use_target_pose, target_pose, progress=gr.Progress()):
124
  # Ensure all dimension numbers are within [0, 512)
125
  num1, num2, num3, num4 = [max(0, min(int(n), 511)) for n in [num1, num2, num3, num4]]
126
 
 
134
  input_data = Image.open(image_input)
135
  input_data = np.array(input_data.resize((112, 112)))
136
 
137
+ generated_images = image_generation(input_data, target_quality, use_target_pose, target_pose, [num1, num2, num3, num4], progress)
138
 
139
  return generated_images
140
 
141
+
142
  def select_image(value, images):
143
  # Convert the float value (0 to 4) to an integer index (0 to 9)
144
+ index = int(value / 2)
145
  return images[index]
146
 
147
  def toggle_inputs(use_pose):
 
164
  <b>Official 🤗 Gradio demo</b> for <a href='https://github.com/HaiyuWu/vec2face' target='_blank'><b>Vec2Face: Scaling Face Dataset Generation with Loosely Constrained Vectors</b></a>.<br>
165
 
166
  How to use:<br>
167
+ 1. Upload an image with a cropped face image or directly click <b>Submit</b> button, three images will be shown on the right.
168
  2. You can control the image quality, image pose, and modify the values in the target dimensions to change the output images.
169
+ 3. The output results will shown three results of dimension modification or pose images.
170
  4. Since the demo is CPU-based, higher quality and larger pose need longer time to run.
171
  5. Enjoy! 😊
172
  """
 
184
 
185
  with gr.Row():
186
  num1 = gr.Number(label="Dimension 1", value=0, minimum=0, maximum=511, step=1)
187
+ num2 = gr.Number(label="Dimension 2", value=50, minimum=0, maximum=511, step=1)
188
+ num3 = gr.Number(label="Dimension 3", value=100, minimum=0, maximum=511, step=1)
189
+ num4 = gr.Number(label="Dimension 4", value=200, minimum=0, maximum=511, step=1)
190
 
191
  random_seed = gr.Number(label="Random Seed", value=42, minimum=0, maximum=MAX_SEED, step=1)
192
  target_quality = gr.Slider(label="Minimum Quality", minimum=22, maximum=35, step=1, value=24)
 
208
 
209
  with gr.Column():
210
  gallery = gr.Image(label="Generated Image")
211
+ generation_time = gr.Textbox(label="Generation Status")
212
  incremental_value_slider = gr.Slider(
213
  label="Result of dimension modification or results of pose images",
214
+ minimum=0, maximum=4, step=2, value=0
215
  )
216
  gr.Markdown("""
217
  - These values are added to the dimensions (before normalization), **please ignore it if pose editing is on**.
 
226
  generated_images = gr.State([])
227
 
228
  submit.click(
229
+ fn=clear_image,
230
+ inputs=[],
231
+ outputs=[gallery]
232
+ ).then(
233
+ fn=generating,
234
+ inputs=[],
235
+ outputs=[generation_time]
236
+ ).then(
237
  fn=process_input,
238
  inputs=[image_file, num1, num2, num3, num4, random_seed, target_quality, use_target_pose, target_pose],
239
  outputs=[generated_images]
240
+ ).then(
241
+ fn=done,
242
+ inputs=[],
243
+ outputs=[generation_time]
244
  ).then(
245
  fn=select_image,
246
  inputs=[incremental_value_slider, generated_images],
247
  outputs=[gallery]
248
  )
249
+ # submit.click(
250
+ # fn=process_input,
251
+ # inputs=[image_file, num1, num2, num3, num4, random_seed, target_quality, use_target_pose, target_pose],
252
+ # outputs=[generated_images]
253
+ # ).then(
254
+ # fn=select_image,
255
+ # inputs=[incremental_value_slider, generated_images],
256
+ # outputs=[gallery]
257
+ # )
258
 
259
  incremental_value_slider.change(
260
  fn=select_image,
 
283
 
284
 
285
  if __name__ == "__main__":
286
+ main()