Update web-demos/hugging_face/app.py

web-demos/hugging_face/app.py

@@ -233,9 +233,33 @@ def show_mask(video_state, interactive_state, mask_dropdown):
 # tracking vos
 def vos_tracking_video(video_state, interactive_state, mask_dropdown):
     operation_log = [("",""), ("Tracking finished! Try to click the Inpainting button to get the inpainting result.","Normal")]
+    model.cutie.clear_memory()
+    if interactive_state["track_end_number"]:
+        following_frames = video_state["origin_images"][video_state["select_frame_number"]:interactive_state["track_end_number"]]
+    else:
+        following_frames = video_state["origin_images"][video_state["select_frame_number"]:]
+
+    if interactive_state["multi_mask"]["masks"]:
+        if len(mask_dropdown) == 0:
+            mask_dropdown = ["mask_001"]
+        mask_dropdown.sort()
+        template_mask = interactive_state["multi_mask"]["masks"][int(mask_dropdown[0].split("_")[1]) - 1] * (int(mask_dropdown[0].split("_")[1]))
+        for i in range(1,len(mask_dropdown)):
+            mask_number = int(mask_dropdown[i].split("_")[1]) - 1 
+            template_mask = np.clip(template_mask+interactive_state["multi_mask"]["masks"][mask_number]*(mask_number+1), 0, mask_number+1)
+        video_state["masks"][video_state["select_frame_number"]]= template_mask
+    else:      
+        template_mask = video_state["masks"][video_state["select_frame_number"]]
+    fps = video_state["fps"]
+
+    # operation error
+    if len(np.unique(template_mask))==1:
+        template_mask[0][0]=1
+        operation_log = [("Please add at least one mask to track by clicking the image in step2.","Error"), ("","")]
+        # return video_output, video_state, interactive_state, operation_error
     masks, logits, painted_images = model.generator(images=following_frames, template_mask=template_mask)
+    # clear GPU memory
     model.cutie.clear_memory()
-    
 
     if interactive_state["track_end_number"]: 
         video_state["masks"][video_state["select_frame_number"]:interactive_state["track_end_number"]] = masks
@@ -251,7 +275,7 @@ def vos_tracking_video(video_state, interactive_state, mask_dropdown):
         orig_h, orig_w = video_state["origin_images"][0].shape[:2]
         for mask in video_state["masks"]:
             mask_up = cv2.resize(mask.astype(np.uint8), (orig_w, orig_h), interpolation=cv2.INTER_NEAREST)
-            binary_mask = np.where(mask_up > 0, 255, 0).astype(np.uint8)
+            binary_mask = np.where(mask > 0, 255, 0).astype(np.uint8)
             bw_frame = np.stack([binary_mask]*3, axis=-1)  # RGB ч/б
             bw_mask_frames.append(bw_frame)
 
@@ -278,8 +302,7 @@ def vos_tracking_video(video_state, interactive_state, mask_dropdown):
         # save_mask(video_state["masks"], video_state["video_name"])
     #### shanggao code for mask save
     return video_output, video_state, interactive_state, operation_log, operation_log
-
-# inpaint 
+    
 def inpaint_video(video_state, *_args):
     operation_log = [("",""), ("Inpainting started in smooth-overlap mode.","Normal")]