Update app.py

app.py

@@ -17,27 +17,14 @@ warnings.filterwarnings('ignore')
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
-# args = option_trans.get_args_parser()
-
 args = option_trans.get_args_parser()
 
 args.dataname = 't2m'
-# args.resume_pth = './output/VQVAE_imp_resnet_100k_hml3d/net_last_VQVAE_res1.pth'
-# args.resume_trans = './output/VQVAE_imp_resnet_100k_hml3d/net_best_fid.pth'
 args.down_t = 2
 args.depth = 3
 args.block_size = 51
 
 vqvae_model = vqvae.HumanVQVAE(args).to(device)
-# ,
-                               # nb_code=512,
-                               # code_dim=512,
-                               # output_emb_width=512,
-                               # down_t=args.down_t,
-                               # stride_t=2,
-                               # width=512,
-                               # depth=args.depth,
-                               # dilation_growth_rate=3).to(device)
 
 transformer_model = trans.Text2Motion_Transformer(num_vq=args.nb_code,
                                                   embed_dim=1024,
@@ -95,7 +82,7 @@ def generate_motion(text, vqvae_model, transformer_model):
     if motion_data.size == 0:
         raise ValueError("Generated motion data is empty")
     return motion_data
-
+    
 def create_animation(joints, title="3D Motion", save_path="static/animation.gif"):
     os.makedirs(os.path.dirname(save_path), exist_ok=True)
     
@@ -107,37 +94,27 @@ def create_animation(joints, title="3D Motion", save_path="static/animation.gif"
     ax = fig.add_subplot(111, projection='3d')
     ax.set_title(title)
 
-    data = np.array(joints).reshape(-1, 22, 3)  # Ensure data shape is (frames, 22, 3)
-    print(f"data shape: {data.shape}")
-
-    # Re-orient the data if necessary
-    data = data[:, :, [0, 2, 1]]  # This swaps Y and Z for example if necessary
-    data[:, :, 2] = -data[:, :, 2]  # Inverting Z if necessary
-
-    # Setting axes limits
-    x_limits = [np.min(data[:, :, 0]), np.max(data[:, :, 0])]
-    y_limits = [np.min(data[:, :, 1]), np.max(data[:, :, 1])]
-    z_limits = [np.min(data[:, :, 2]), np.max(data[:, :, 2])]
-    ax.set_xlim(x_limits)
-    ax.set_ylim(y_limits)
-    ax.set_zlim(z_limits)
-
-    # Preparing lines to plot based on connections
-    lines = []
-    for conn in connections:
-        x_coords = data[0, [conn[0], conn[1]], 0]
-        y_coords = data[0, [conn[0], conn[1]], 1]
-        z_coords = data[0, [conn[0], conn[1]], 2]
-        line, = ax.plot(x_coords, y_coords, z_coords, marker='o', markersize=5)
-        lines.append(line)
+    data = np.array(joints).reshape(-1, 22, 3)
+    data[:, :, [1, 2]] = data[:, :, [2, 1]] 
+    data[:, :, 2] = -data[:, :, 2]
+
+    x_min, x_max = data[:, :, 0].min(), data[:, :, 0].max()
+    y_min, y_max = data[:, :, 1].min(), data[:, :, 1].max()
+    z_min, z_max = data[:, :, 2].min(), data[:, :, 2].max()
+    ax.set_xlim([x_min, x_max])
+    ax.set_ylim([y_min, y_max])
+    ax.set_zlim([z_min, z_max])
+
+    lines = [ax.plot([data[0, conn[0], 0], data[0, conn[1], 0]],
+                     [data[0, conn[0], 1], data[0, conn[1], 1]],
+                     [data[0, conn[0], 2], data[0, conn[1], 2]], 
+                     marker='o', markersize=5)[0] for conn in connections]
 
     def update(num):
         for line, conn in zip(lines, connections):
-            x = data[num, [conn[0], conn[1]], 0]
-            y = data[num, [conn[0], conn[1]], 1]
-            z = data[num, [conn[0], conn[1]], 2]
-            line.set_data(x, y)
-            line.set_3d_properties(z)
+            line.set_data([data[num, conn[0], 0], data[num, conn[1], 0]], 
+                          [data[num, conn[0], 1], data[num, conn[1], 1]])
+            line.set_3d_properties([data[num, conn[0], 2], data[num, conn[1], 2]])
         return lines
 
     ani = FuncAnimation(fig, update, frames=len(data), interval=50, blit=True)
@@ -145,7 +122,6 @@ def create_animation(joints, title="3D Motion", save_path="static/animation.gif"
     plt.close(fig)
     return save_path
 
-
 def infer(text):
     motion_data = generate_motion(text, vqvae_model, transformer_model)
     if motion_data.size == 0:
@@ -155,8 +131,8 @@ def infer(text):
 
 with gr.Blocks(css=".container { max-width: 800px; margin: auto; }") as demo:
     with gr.Column():
-        gr.Markdown("## 3D Human Motion Generation on " + ("GPU" if device == "cuda" else "CPU"))
-        text_input = gr.Textbox(label="Describe the action", placeholder="Enter text description for the action here...")
+        gr.Markdown("## 3D Human Motion" + ("Generation" if device == "cuda" else "CPU"))
+        text_input = gr.Textbox(label="Enter the human action to generate", placeholder="Enter text description for the action here...")
         output_image = gr.Image(label="Generated Human Motion")
         submit_button = gr.Button("Generate Motion")