33 lines
1.3 KiB
Python
33 lines
1.3 KiB
Python
import os
|
|
import torch
|
|
#import argparse
|
|
import numpy as np
|
|
from PIL import Image
|
|
import random
|
|
import imageio
|
|
import folder_paths
|
|
|
|
|
|
def get_mask(video_list,fps,local_sam, x=255, y=255):
|
|
from ultralytics.models.sam import SAM2VideoPredictor #ultralytics>=8.3.0
|
|
overrides = dict(conf=0.25, task="segment", mode="predict", imgsz=1024, model=local_sam)
|
|
predictor = SAM2VideoPredictor(overrides=overrides)
|
|
file_prefix = ''.join(random.choice("0123456789") for _ in range(6))
|
|
frames = [np.array(img) for img in video_list]
|
|
video_file = os.path.join(folder_paths.get_input_directory(), f"audio_{file_prefix}_temp.mp4")
|
|
imageio.mimsave(video_file, frames, fps=fps, codec='libx264')
|
|
results = predictor(source=video_file,points=[x, y],labels=[1])
|
|
mask_list=[]
|
|
for i in range(len(results)):
|
|
mask = (results[i].masks.data).squeeze().to(torch.float16)
|
|
mask = (mask * 255).cpu().numpy().astype(np.uint8)
|
|
mask_image = Image.fromarray(mask).convert('RGB')
|
|
mask_list.append(mask_image)
|
|
# mask_dir = f'masks_animatediff/{video_name}'
|
|
# if not os.path.exists(mask_dir):
|
|
# os.makedirs(mask_dir)
|
|
mask_image.save(folder_paths.get_input_directory() + f'/{str(i).zfill(3)}.png')
|
|
return mask_list
|
|
|
|
|