Update app.py

app.py

@@ -527,6 +527,154 @@ class Drag:
         return val_save_dir
 
 
+def reset_states(first_frame_path, last_frame_path, tracking_points):
+    first_frame_path = gr.State()
+    last_frame_path = gr.State()
+    tracking_points = gr.State([])
+
+    return first_frame_path, last_frame_path, tracking_points
+
+
+def preprocess_image(image):
+
+    image_pil = image2pil(image.name)
+
+    raw_w, raw_h = image_pil.size
+    # resize_ratio = max(512 / raw_w, 320 / raw_h)
+    # image_pil = image_pil.resize((int(raw_w * resize_ratio), int(raw_h * resize_ratio)), Image.BILINEAR)
+    # image_pil = transforms.CenterCrop((320, 512))(image_pil.convert('RGB'))
+    image_pil = image_pil.resize((512, 320), Image.BILINEAR)
+
+    first_frame_path = os.path.join(args.output_dir, f"first_frame_{str(uuid.uuid4())[:4]}.png")
+    
+    image_pil.save(first_frame_path)
+
+    return first_frame_path, first_frame_path, gr.State([])
+
+
+def preprocess_image_end(image_end):
+
+    image_end_pil = image2pil(image_end.name)
+
+    raw_w, raw_h = image_end_pil.size
+    # resize_ratio = max(512 / raw_w, 320 / raw_h)
+    # image_end_pil = image_end_pil.resize((int(raw_w * resize_ratio), int(raw_h * resize_ratio)), Image.BILINEAR)
+    # image_end_pil = transforms.CenterCrop((320, 512))(image_end_pil.convert('RGB'))
+    image_end_pil = image_end_pil.resize((512, 320), Image.BILINEAR)
+
+    last_frame_path = os.path.join(args.output_dir, f"last_frame_{str(uuid.uuid4())[:4]}.png")
+
+    image_end_pil.save(last_frame_path)
+
+    return last_frame_path, last_frame_path, gr.State([])
+
+
+def add_drag(tracking_points):
+    tracking_points.constructor_args['value'].append([])
+    return tracking_points
+
+
+def delete_last_drag(tracking_points, first_frame_path, last_frame_path):
+    tracking_points.constructor_args['value'].pop()
+    transparent_background = Image.open(first_frame_path).convert('RGBA')
+    transparent_background_end = Image.open(last_frame_path).convert('RGBA')
+    w, h = transparent_background.size
+    transparent_layer = np.zeros((h, w, 4))
+
+    for track in tracking_points.constructor_args['value']:
+        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), (255, 0, 0, 255), 2, tipLength=8 / arrow_length)
+                else:
+                    cv2.line(transparent_layer, tuple(start_point), tuple(end_point), (255, 0, 0, 255), 2,)
+        else:
+            cv2.circle(transparent_layer, tuple(track[0]), 5, (255, 0, 0, 255), -1)
+
+    transparent_layer = Image.fromarray(transparent_layer.astype(np.uint8))
+    trajectory_map = Image.alpha_composite(transparent_background, transparent_layer)
+    trajectory_map_end = Image.alpha_composite(transparent_background_end, transparent_layer)
+
+    return tracking_points, trajectory_map, trajectory_map_end
+
+
+def delete_last_step(tracking_points, first_frame_path, last_frame_path):
+    tracking_points.constructor_args['value'][-1].pop()
+    transparent_background = Image.open(first_frame_path).convert('RGBA')
+    transparent_background_end = Image.open(last_frame_path).convert('RGBA')
+    w, h = transparent_background.size
+    transparent_layer = np.zeros((h, w, 4))
+
+    for track in tracking_points.constructor_args['value']:
+        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), (255, 0, 0, 255), 2, tipLength=8 / arrow_length)
+                else:
+                    cv2.line(transparent_layer, tuple(start_point), tuple(end_point), (255, 0, 0, 255), 2,)
+        else:
+            cv2.circle(transparent_layer, tuple(track[0]), 5, (255, 0, 0, 255), -1)
+
+    transparent_layer = Image.fromarray(transparent_layer.astype(np.uint8))
+    trajectory_map = Image.alpha_composite(transparent_background, transparent_layer)
+    trajectory_map_end = Image.alpha_composite(transparent_background_end, transparent_layer)
+
+    return tracking_points, trajectory_map, trajectory_map_end
+
+
+def add_tracking_points(tracking_points, first_frame_path, last_frame_path, evt: gr.SelectData):  # SelectData is a subclass of EventData
+    print(f"You selected {evt.value} at {evt.index} from {evt.target}")
+    tracking_points.constructor_args['value'][-1].append(evt.index)
+
+    transparent_background = Image.open(first_frame_path).convert('RGBA')
+    transparent_background_end = Image.open(last_frame_path).convert('RGBA')
+
+    w, h = transparent_background.size
+    transparent_layer = 0
+    for idx, track in enumerate(tracking_points.constructor_args['value']):
+        # mask = cv2.imread(
+        #     os.path.join(args.output_dir, f"mask_{idx+1}.jpg")
+        # )
+        mask = np.zeros((320, 512, 3))
+        color = color_list[idx+1]
+        transparent_layer = mask[:, :, 0].reshape(h, w, 1) * color.reshape(1, 1, -1) + transparent_layer
+
+        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), (255, 0, 0, 255), 2, tipLength=8 / arrow_length)
+                else:
+                    cv2.line(transparent_layer, tuple(start_point), tuple(end_point), (255, 0, 0, 255), 2,)
+        else:
+            cv2.circle(transparent_layer, tuple(track[0]), 5, (255, 0, 0, 255), -1)
+
+    transparent_layer = Image.fromarray(transparent_layer.astype(np.uint8))
+    alpha_coef = 0.99
+    im2_data = transparent_layer.getdata()
+    new_im2_data = [(r, g, b, int(a * alpha_coef)) for r, g, b, a in im2_data]
+    transparent_layer.putdata(new_im2_data)
+
+    trajectory_map = Image.alpha_composite(transparent_background, transparent_layer)
+    trajectory_map_end = Image.alpha_composite(transparent_background_end, transparent_layer)
+
+    return tracking_points, trajectory_map, trajectory_map_end
+
+
 if __name__ == "__main__":
 
     args = get_args()
@@ -564,153 +712,6 @@ if __name__ == "__main__":
         last_frame_path = gr.State()
         tracking_points = gr.State([])
     
-        def reset_states(first_frame_path, last_frame_path, tracking_points):
-            first_frame_path = gr.State()
-            last_frame_path = gr.State()
-            tracking_points = gr.State([])
-    
-            return first_frame_path, last_frame_path, tracking_points
-    
-    
-        def preprocess_image(image):
-    
-            image_pil = image2pil(image.name)
-    
-            raw_w, raw_h = image_pil.size
-            # resize_ratio = max(512 / raw_w, 320 / raw_h)
-            # image_pil = image_pil.resize((int(raw_w * resize_ratio), int(raw_h * resize_ratio)), Image.BILINEAR)
-            # image_pil = transforms.CenterCrop((320, 512))(image_pil.convert('RGB'))
-            image_pil = image_pil.resize((512, 320), Image.BILINEAR)
-    
-            first_frame_path = os.path.join(args.output_dir, f"first_frame_{str(uuid.uuid4())[:4]}.png")
-            
-            image_pil.save(first_frame_path)
-    
-            return first_frame_path, first_frame_path, gr.State([])
-    
-    
-        def preprocess_image_end(image_end):
-    
-            image_end_pil = image2pil(image_end.name)
-    
-            raw_w, raw_h = image_end_pil.size
-            # resize_ratio = max(512 / raw_w, 320 / raw_h)
-            # image_end_pil = image_end_pil.resize((int(raw_w * resize_ratio), int(raw_h * resize_ratio)), Image.BILINEAR)
-            # image_end_pil = transforms.CenterCrop((320, 512))(image_end_pil.convert('RGB'))
-            image_end_pil = image_end_pil.resize((512, 320), Image.BILINEAR)
-    
-            last_frame_path = os.path.join(args.output_dir, f"last_frame_{str(uuid.uuid4())[:4]}.png")
-    
-            image_end_pil.save(last_frame_path)
-    
-            return last_frame_path, last_frame_path, gr.State([])
-    
-    
-        def add_drag(tracking_points):
-            tracking_points.constructor_args['value'].append([])
-            return tracking_points
-    
-    
-        def delete_last_drag(tracking_points, first_frame_path, last_frame_path):
-            tracking_points.constructor_args['value'].pop()
-            transparent_background = Image.open(first_frame_path).convert('RGBA')
-            transparent_background_end = Image.open(last_frame_path).convert('RGBA')
-            w, h = transparent_background.size
-            transparent_layer = np.zeros((h, w, 4))
-    
-            for track in tracking_points.constructor_args['value']:
-                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), (255, 0, 0, 255), 2, tipLength=8 / arrow_length)
-                        else:
-                            cv2.line(transparent_layer, tuple(start_point), tuple(end_point), (255, 0, 0, 255), 2,)
-                else:
-                    cv2.circle(transparent_layer, tuple(track[0]), 5, (255, 0, 0, 255), -1)
-    
-            transparent_layer = Image.fromarray(transparent_layer.astype(np.uint8))
-            trajectory_map = Image.alpha_composite(transparent_background, transparent_layer)
-            trajectory_map_end = Image.alpha_composite(transparent_background_end, transparent_layer)
-    
-            return tracking_points, trajectory_map, trajectory_map_end
-    
-    
-        def delete_last_step(tracking_points, first_frame_path, last_frame_path):
-            tracking_points.constructor_args['value'][-1].pop()
-            transparent_background = Image.open(first_frame_path).convert('RGBA')
-            transparent_background_end = Image.open(last_frame_path).convert('RGBA')
-            w, h = transparent_background.size
-            transparent_layer = np.zeros((h, w, 4))
-    
-            for track in tracking_points.constructor_args['value']:
-                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), (255, 0, 0, 255), 2, tipLength=8 / arrow_length)
-                        else:
-                            cv2.line(transparent_layer, tuple(start_point), tuple(end_point), (255, 0, 0, 255), 2,)
-                else:
-                    cv2.circle(transparent_layer, tuple(track[0]), 5, (255, 0, 0, 255), -1)
-    
-            transparent_layer = Image.fromarray(transparent_layer.astype(np.uint8))
-            trajectory_map = Image.alpha_composite(transparent_background, transparent_layer)
-            trajectory_map_end = Image.alpha_composite(transparent_background_end, transparent_layer)
-    
-            return tracking_points, trajectory_map, trajectory_map_end
-    
-    
-        def add_tracking_points(tracking_points, first_frame_path, last_frame_path, evt: gr.SelectData):  # SelectData is a subclass of EventData
-            print(f"You selected {evt.value} at {evt.index} from {evt.target}")
-            tracking_points.constructor_args['value'][-1].append(evt.index)
-    
-            transparent_background = Image.open(first_frame_path).convert('RGBA')
-            transparent_background_end = Image.open(last_frame_path).convert('RGBA')
-    
-            w, h = transparent_background.size
-            transparent_layer = 0
-            for idx, track in enumerate(tracking_points.constructor_args['value']):
-                # mask = cv2.imread(
-                #     os.path.join(args.output_dir, f"mask_{idx+1}.jpg")
-                # )
-                mask = np.zeros((320, 512, 3))
-                color = color_list[idx+1]
-                transparent_layer = mask[:, :, 0].reshape(h, w, 1) * color.reshape(1, 1, -1) + transparent_layer
-    
-                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), (255, 0, 0, 255), 2, tipLength=8 / arrow_length)
-                        else:
-                            cv2.line(transparent_layer, tuple(start_point), tuple(end_point), (255, 0, 0, 255), 2,)
-                else:
-                    cv2.circle(transparent_layer, tuple(track[0]), 5, (255, 0, 0, 255), -1)
-    
-            transparent_layer = Image.fromarray(transparent_layer.astype(np.uint8))
-            alpha_coef = 0.99
-            im2_data = transparent_layer.getdata()
-            new_im2_data = [(r, g, b, int(a * alpha_coef)) for r, g, b, a in im2_data]
-            transparent_layer.putdata(new_im2_data)
-    
-            trajectory_map = Image.alpha_composite(transparent_background, transparent_layer)
-            trajectory_map_end = Image.alpha_composite(transparent_background_end, transparent_layer)
-    
-            return tracking_points, trajectory_map, trajectory_map_end
-    
         with gr.Row():
             with gr.Column(scale=1):
                 image_upload_button = gr.UploadButton(label="Upload Start Image", file_types=["image"])
@@ -798,4 +799,4 @@ if __name__ == "__main__":
     
         run_button.click(Framer.run, [first_frame_path, last_frame_path, tracking_points, controlnet_cond_scale, motion_bucket_id], output_video)
     
-        demo.launch(max_threads=1)
+    demo.launch()