Update web-demos/hugging_face/app.py

web-demos/hugging_face/app.py

@@ -13,7 +13,6 @@ import torch
 import torchvision
 import numpy as np
 import gradio as gr
-from PIL import Image
 
 from tools.painter import mask_painter
 from track_anything import TrackingAnything
@@ -253,14 +252,10 @@ def vos_tracking_video(video_state, interactive_state, mask_dropdown):
         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, alpha_visuals = model.generator(images=following_frames, template_mask=template_mask)
-
+    masks, logits, painted_images = model.generator(images=following_frames, template_mask=template_mask)
     # clear GPU memory
     model.cutie.clear_memory()
 
-    # сохранить альфа-канальные маски в состояние (для отображения или сохранения видео)
-    video_state["alpha_visuals"] = alpha_visuals
-
     if interactive_state["track_end_number"]: 
         video_state["masks"][video_state["select_frame_number"]:interactive_state["track_end_number"]] = masks
         video_state["logits"][video_state["select_frame_number"]:interactive_state["track_end_number"]] = logits
@@ -272,10 +267,6 @@ def vos_tracking_video(video_state, interactive_state, mask_dropdown):
 
     video_output = generate_video_from_frames(video_state["painted_images"], output_path="./result/track/{}".format(video_state["video_name"]), fps=float(fps)) # import video_input to name the output video
     interactive_state["inference_times"] += 1
-    # Дополнительно: альфа-маска-видео
-    if "alpha_visuals" in video_state:
-        generate_video_from_frames(video_state["alpha_visuals"], output_path="./result/track/{}_alpha.mp4".format(video_state["video_name"].split('.')[0]), fps=float(fps), is_rgba=True)
-
 
     print("Tracking resolution:", following_frames[0].shape)
     
@@ -334,7 +325,7 @@ def inpaint_video(video_state, resize_ratio_number, dilate_radius_number, raft_i
 
 
 # generate video after vos inference
-def generate_video_from_frames(frames, output_path, fps=30,is_rgba=False):
+def generate_video_from_frames(frames, output_path, fps=30):
     """
     Generates a video from a list of frames.
     
@@ -365,11 +356,6 @@ def generate_video_from_frames(frames, output_path, fps=30,is_rgba=False):
     if not os.path.exists(os.path.dirname(output_path)):
         os.makedirs(os.path.dirname(output_path))
 
-    if is_rgba:
-        frames = torch.from_numpy(np.asarray(frames).astype(np.uint8))
-    else:
-        frames = torch.from_numpy(np.asarray(frames))
-
     # Write the video
     torchvision.io.write_video(output_path, frames, fps=fps, video_codec="libx264")
     return output_path