update

app.py

@@ -10,6 +10,23 @@ from sixdrepnet.model import SixDRepNet
 import pixel_generator.vec2face.model_vec2face as model_vec2face
 MAX_SEED = np.iinfo(np.int32).max
 import torch
+from time import time
+
+
+def clear_image():
+    return None
+
+
+def clear_generation_time():
+    return ""
+
+
+def generating():
+    return "Generating images..."
+
+
+def done():
+    return "Done!"
 
 
 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():
     return generator, id_model, pose_model, quality_model
 
 
-def image_generation(input_image, quality, use_target_pose, pose, dimension):
+def image_generation(input_image, quality, use_target_pose, pose, dimension, progress=gr.Progress()):
     generator, id_model, pose_model, quality_model = initialize_models()
 
     generated_images = []
@@ -77,12 +94,10 @@ def image_generation(input_image, quality, use_target_pose, pose, dimension):
     if not use_target_pose:
         features = []
         norm = np.linalg.norm(feature, 2, 1, True)
-        for i in np.arange(0, 4.8, 0.8):
+        for i in progress.tqdm(np.arange(0, 4.8, 2), desc="Generating images"):
             updated_feature = feature
             updated_feature[0][dimension] = feature[0][dimension] + i
-
             updated_feature = updated_feature / np.linalg.norm(updated_feature, 2, 1, True) * norm
-
             features.append(updated_feature)
         features = torch.tensor(np.vstack(features)).float()
         if quality > 25:
@@ -90,7 +105,7 @@ def image_generation(input_image, quality, use_target_pose, pose, dimension):
         else:
             _, _, images, *_ = generator(features)
     else:
-        features = torch.repeat_interleave(torch.tensor(feature), 6, dim=0)
+        features = torch.repeat_interleave(torch.tensor(feature), 3, dim=0)
         features = sample_nearby_vectors(features, [0.7], [1]).float()
         if quality > 25 or pose > 20:
             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):
             _, _, images, *_ = generator(features)
 
     images = ((images.permute(0, 2, 3, 1).detach().cpu().numpy() + 1) / 2 * 255).astype(np.uint8)
-    for image in images:
+    for image in progress.tqdm(images, desc="Processing images"):
         generated_images.append(Image.fromarray(image))
+
     return generated_images
 
 
-def process_input(image_input, num1, num2, num3, num4, random_seed, target_quality, use_target_pose, target_pose):
+def process_input(image_input, num1, num2, num3, num4, random_seed, target_quality, use_target_pose, target_pose, progress=gr.Progress()):
     # Ensure all dimension numbers are within [0, 512)
     num1, num2, num3, num4 = [max(0, min(int(n), 511)) for n in [num1, num2, num3, num4]]
 
@@ -118,13 +134,14 @@ def process_input(image_input, num1, num2, num3, num4, random_seed, target_quali
         input_data = Image.open(image_input)
         input_data = np.array(input_data.resize((112, 112)))
 
-    generated_images = image_generation(input_data, target_quality, use_target_pose, target_pose, [num1, num2, num3, num4])
+    generated_images = image_generation(input_data, target_quality, use_target_pose, target_pose, [num1, num2, num3, num4], progress)
 
     return generated_images
 
+
 def select_image(value, images):
     # Convert the float value (0 to 4) to an integer index (0 to 9)
-    index = int(value / 0.8)
+    index = int(value / 2)
     return images[index]
 
 def toggle_inputs(use_pose):
@@ -147,9 +164,9 @@ def main():
             <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>
 
             How to use:<br>
-            1. Upload an image with a cropped face image or directly click <b>Submit</b> button, six images will be shown on the right. 
+            1. Upload an image with a cropped face image or directly click <b>Submit</b> button, three images will be shown on the right. 
             2. You can control the image quality, image pose, and modify the values in the target dimensions to change the output images. 
-            3. The output results will shown six results of dimension modification or pose images.
+            3. The output results will shown three results of dimension modification or pose images.
             4. Since the demo is CPU-based, higher quality and larger pose need longer time to run.
             5. Enjoy! 😊
             """
@@ -167,9 +184,9 @@ def main():
 
                 with gr.Row():
                     num1 = gr.Number(label="Dimension 1", value=0, minimum=0, maximum=511, step=1)
-                    num2 = gr.Number(label="Dimension 2", value=0, minimum=0, maximum=511, step=1)
-                    num3 = gr.Number(label="Dimension 3", value=0, minimum=0, maximum=511, step=1)
-                    num4 = gr.Number(label="Dimension 4", value=0, minimum=0, maximum=511, step=1)
+                    num2 = gr.Number(label="Dimension 2", value=50, minimum=0, maximum=511, step=1)
+                    num3 = gr.Number(label="Dimension 3", value=100, minimum=0, maximum=511, step=1)
+                    num4 = gr.Number(label="Dimension 4", value=200, minimum=0, maximum=511, step=1)
 
                 random_seed = gr.Number(label="Random Seed", value=42, minimum=0, maximum=MAX_SEED, step=1)
                 target_quality = gr.Slider(label="Minimum Quality", minimum=22, maximum=35, step=1, value=24)
@@ -191,9 +208,10 @@ def main():
 
             with gr.Column():
                 gallery = gr.Image(label="Generated Image")
+                generation_time = gr.Textbox(label="Generation Status")
                 incremental_value_slider = gr.Slider(
                     label="Result of dimension modification or results of pose images",
-                    minimum=0, maximum=4, step=0.8, value=0
+                    minimum=0, maximum=4, step=2, value=0
                 )
                 gr.Markdown("""
                             - These values are added to the dimensions (before normalization), **please ignore it if pose editing is on**.
@@ -208,14 +226,35 @@ def main():
         generated_images = gr.State([])
 
         submit.click(
+            fn=clear_image,
+            inputs=[],
+            outputs=[gallery]
+        ).then(
+            fn=generating,
+            inputs=[],
+            outputs=[generation_time]
+        ).then(
             fn=process_input,
             inputs=[image_file, num1, num2, num3, num4, random_seed, target_quality, use_target_pose, target_pose],
             outputs=[generated_images]
+        ).then(
+            fn=done,
+            inputs=[],
+            outputs=[generation_time]
         ).then(
             fn=select_image,
             inputs=[incremental_value_slider, generated_images],
             outputs=[gallery]
         )
+        # submit.click(
+        #     fn=process_input,
+        #     inputs=[image_file, num1, num2, num3, num4, random_seed, target_quality, use_target_pose, target_pose],
+        #     outputs=[generated_images]
+        # ).then(
+        #     fn=select_image,
+        #     inputs=[incremental_value_slider, generated_images],
+        #     outputs=[gallery]
+        # )
 
         incremental_value_slider.change(
             fn=select_image,
@@ -244,4 +283,4 @@ def main():
 
 
 if __name__ == "__main__":
-    main()
+    main()