Process examples when loaded

app.py

@@ -297,23 +297,11 @@ class ImageConductor:
         guidance_scale,
         num_inference_steps,
         personalized,
-        examples_type,
     ):
         print("Run!")
-        if examples_type != "":
-            ### for adapting high version gradio
-            tracking_points = gr.State([])
-            first_frame_path = IMAGE_PATH[examples_type]
-            points = json.load(open(POINTS[examples_type]))
-            tracking_points.value.extend(points)
-            print("example first_frame_path", first_frame_path)
-            print("example tracking_points", tracking_points.value)
 
         original_width, original_height = 384, 256
-        if isinstance(tracking_points, list):
-            input_all_points = tracking_points
-        else:
-            input_all_points = tracking_points.value
+        input_all_points = tracking_points
 
         print("input_all_points", input_all_points)
         resized_all_points = [
@@ -415,7 +403,7 @@ class ImageConductor:
         # outputs_path = os.path.join(output_dir, f'output_{i}_{id}.gif')
         # save_videos_grid(sample[0][None], outputs_path)
         print("Done!")
-        return {output_image: visualized_drag, output_video: outputs_path}
+        return visualized_drag, outputs_path
 
 
 def reset_states(first_frame_path, tracking_points):
@@ -487,6 +475,54 @@ def add_tracking_points(
     return {tracking_points_var: tracking_points, input_image: trajectory_map}
 
 
+def preprocess_example_image(image_path, tracking_points, drag_mode):
+    image_pil = image2pil(image_path)
+    raw_w, raw_h = image_pil.size
+    resize_ratio = max(384 / raw_w, 256 / raw_h)
+    image_pil = image_pil.resize((int(raw_w * resize_ratio), int(raw_h * resize_ratio)), Image.BILINEAR)
+    image_pil = transforms.CenterCrop((256, 384))(image_pil.convert("RGB"))
+    id = str(uuid.uuid4())[:4]
+    first_frame_path = os.path.join(output_dir, f"first_frame_{id}.jpg")
+    image_pil.save(first_frame_path, quality=95)
+
+    if drag_mode == "object":
+        color = (255, 0, 0, 255)
+    elif drag_mode == "camera":
+        color = (0, 0, 255, 255)
+
+    transparent_background = Image.open(first_frame_path).convert("RGBA")
+    w, h = transparent_background.size
+    transparent_layer = np.zeros((h, w, 4))
+
+    for track in tracking_points:
+        if len(track) > 1:
+            for i in range(len(track) - 1):
+                start_point = track[i]
+                end_point = track[i + 1]
+                vx = end_point[0] - start_point[0]
+                vy = end_point[1] - start_point[1]
+                arrow_length = np.sqrt(vx**2 + vy**2)
+                if i == len(track) - 2:
+                    cv2.arrowedLine(
+                        transparent_layer, tuple(start_point), tuple(end_point), color, 2, tipLength=8 / arrow_length
+                    )
+                else:
+                    cv2.line(
+                        transparent_layer,
+                        tuple(start_point),
+                        tuple(end_point),
+                        color,
+                        2,
+                    )
+        else:
+            cv2.circle(transparent_layer, tuple(track[0]), 5, color, -1)
+
+    transparent_layer = Image.fromarray(transparent_layer.astype(np.uint8))
+    trajectory_map = Image.alpha_composite(transparent_background, transparent_layer)
+
+    return trajectory_map, first_frame_path
+
+
 def add_drag(tracking_points):
     if not tracking_points or tracking_points[-1]:
         tracking_points.append([])
@@ -571,6 +607,15 @@ def delete_last_step(tracking_points, first_frame_path, drag_mode):
     return {tracking_points_var: tracking_points, input_image: trajectory_map}
 
 
+def load_example(drag_mode, examples_type):
+    example_image_path = IMAGE_PATH[examples_type]
+    with open(POINTS[examples_type]) as f:
+        tracking_points = json.load(f)
+    tracking_points = np.round(tracking_points).astype(int).tolist()
+    trajectory_map, first_frame_path = preprocess_example_image(example_image_path, tracking_points, drag_mode)
+    return {input_image: trajectory_map, first_frame_path_var: first_frame_path, tracking_points_var: tracking_points}
+
+
 device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
 ImageConductor_net = ImageConductor(
     device=device,
@@ -725,9 +770,16 @@ with block:
             guidance_scale,
             num_inference_steps,
             personalized,
-            examples_type,
         ],
         [output_image, output_video],
     )
 
+    examples_type.change(
+        fn=load_example,
+        inputs=[drag_mode, examples_type],
+        outputs=[input_image, first_frame_path_var, tracking_points_var],
+        api_name=False,
+        queue=False,
+    )
+
 block.queue().launch()