Update web-demos/hugging_face/app.py

web-demos/hugging_face/app.py

@@ -13,6 +13,7 @@ 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
@@ -252,10 +253,14 @@ 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 = model.generator(images=following_frames, template_mask=template_mask)
+    masks, logits, painted_images, alpha_visuals = 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
@@ -267,6 +272,10 @@ 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)
     
@@ -325,7 +334,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):
+def generate_video_from_frames(frames, output_path, fps=30,is_rgba=False):
     """
     Generates a video from a list of frames.
     
@@ -356,6 +365,11 @@ def generate_video_from_frames(frames, output_path, fps=30):
     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