import glob import logging import os from pathlib import Path import cv2 import numpy as np import torch from groundingdino.models import build_model from groundingdino.util.slconfig import SLConfig from groundingdino.util.utils import clean_state_dict, get_phrases_from_posmap from PIL import Image from segment_anything_hq import (SamPredictor, build_sam_vit_b, build_sam_vit_h, build_sam_vit_l) from segment_anything_hq.build_sam import build_sam_vit_t from tqdm.rich import tqdm logger = logging.getLogger(__name__) build_sam_table={ "sam_hq_vit_l":build_sam_vit_l, "sam_hq_vit_h":build_sam_vit_h, "sam_hq_vit_b":build_sam_vit_b, "sam_hq_vit_tiny":build_sam_vit_t, } # adapted from https://github.com/IDEA-Research/Grounded-Segment-Anything/blob/main/grounded_sam_demo.py class MaskPredictor: def __init__(self,model_config_path, model_checkpoint_path,device, sam_checkpoint, box_threshold=0.3, text_threshold=0.25 ): self.groundingdino_model = None self.sam_predictor = None self.model_config_path = model_config_path self.model_checkpoint_path = model_checkpoint_path self.device = device self.sam_checkpoint = sam_checkpoint self.box_threshold = box_threshold self.text_threshold = text_threshold def load_groundingdino_model(self): args = SLConfig.fromfile(self.model_config_path) args.device = self.device model = build_model(args) checkpoint = torch.load(self.model_checkpoint_path, map_location="cpu") load_res = model.load_state_dict(clean_state_dict(checkpoint["model"]), strict=False) #print(load_res) _ = model.eval() self.groundingdino_model = model def load_sam_predictor(self): s = Path(self.sam_checkpoint) self.sam_predictor = SamPredictor(build_sam_table[ s.stem ](checkpoint=self.sam_checkpoint).to(self.device)) def transform_image(self,image_pil): import groundingdino.datasets.transforms as T transform = T.Compose( [ T.RandomResize([800], max_size=1333), T.ToTensor(), T.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]), ] ) image, _ = transform(image_pil, None) # 3, h, w return image def get_grounding_output(self, image, caption, with_logits=True): model = self.groundingdino_model device = self.device caption = caption.lower() caption = caption.strip() if not caption.endswith("."): caption = caption + "." model = model.to(device) image = image.to(device) with torch.no_grad(): outputs = model(image[None], captions=[caption]) logits = outputs["pred_logits"].cpu().sigmoid()[0] # (nq, 256) boxes = outputs["pred_boxes"].cpu()[0] # (nq, 4) logits.shape[0] # filter output logits_filt = logits.clone() boxes_filt = boxes.clone() filt_mask = logits_filt.max(dim=1)[0] > self.box_threshold logits_filt = logits_filt[filt_mask] # num_filt, 256 boxes_filt = boxes_filt[filt_mask] # num_filt, 4 logits_filt.shape[0] # get phrase tokenlizer = model.tokenizer tokenized = tokenlizer(caption) # build pred pred_phrases = [] for logit, box in zip(logits_filt, boxes_filt): pred_phrase = get_phrases_from_posmap(logit > self.text_threshold, tokenized, tokenlizer) if with_logits: pred_phrases.append(pred_phrase + f"({str(logit.max().item())[:4]})") else: pred_phrases.append(pred_phrase) return boxes_filt, pred_phrases def __call__(self, image_pil:Image, text_prompt): if self.groundingdino_model is None: self.load_groundingdino_model() self.load_sam_predictor() transformed_img = self.transform_image(image_pil) # run grounding dino model boxes_filt, pred_phrases = self.get_grounding_output( transformed_img, text_prompt ) if boxes_filt.shape[0] == 0: logger.info(f"object not found") w, h = image_pil.size return np.zeros(shape=(1,h,w), dtype=bool) img_array = np.array(image_pil) self.sam_predictor.set_image(img_array) size = image_pil.size H, W = size[1], size[0] for i in range(boxes_filt.size(0)): boxes_filt[i] = boxes_filt[i] * torch.Tensor([W, H, W, H]) boxes_filt[i][:2] -= boxes_filt[i][2:] / 2 boxes_filt[i][2:] += boxes_filt[i][:2] boxes_filt = boxes_filt.cpu() transformed_boxes = self.sam_predictor.transform.apply_boxes_torch(boxes_filt, img_array.shape[:2]).to(self.device) masks, _, _ = self.sam_predictor.predict_torch( point_coords = None, point_labels = None, boxes = transformed_boxes.to(self.device), multimask_output = False, ) result = None for m in masks: if result is None: result = m else: result |= m result = result.cpu().detach().numpy().copy() return result def load_mask_list(mask_dir, masked_area_list, mask_padding): mask_frame_list = sorted(glob.glob( os.path.join(mask_dir, "[0-9]*.png"), recursive=False)) kernel = np.ones((abs(mask_padding),abs(mask_padding)),np.uint8) for m in mask_frame_list: cur = int(Path(m).stem) tmp = np.asarray(Image.open(m)) if mask_padding < 0: tmp = cv2.erode(tmp, kernel,iterations = 1) elif mask_padding > 0: tmp = cv2.dilate(tmp, kernel,iterations = 1) masked_area_list[cur] = tmp[None,...] return masked_area_list def crop_mask_list(mask_list): area_list = [] max_h = 0 max_w = 0 for m in mask_list: if m is None: area_list.append(None) continue m = m > 127 area = np.where(m[0] == True) if area[0].size == 0: area_list.append(None) continue ymin = min(area[0]) ymax = max(area[0]) xmin = min(area[1]) xmax = max(area[1]) h = ymax+1 - ymin w = xmax+1 - xmin max_h = max(max_h, h) max_w = max(max_w, w) area_list.append( (ymin, ymax, xmin, xmax) ) #crop = m[ymin:ymax+1,xmin:xmax+1] logger.info(f"{max_h=}") logger.info(f"{max_w=}") border_h = mask_list[0].shape[1] border_w = mask_list[0].shape[2] mask_pos_list=[] cropped_mask_list=[] for a, m in zip(area_list, mask_list): if m is None or a is None: mask_pos_list.append(None) cropped_mask_list.append(None) continue ymin,ymax,xmin,xmax = a h = ymax+1 - ymin w = xmax+1 - xmin # H diff_h = max_h - h dh1 = diff_h//2 dh2 = diff_h - dh1 y1 = ymin - dh1 y2 = ymax + dh2 if y1 < 0: y1 = 0 y2 = max_h-1 elif y2 >= border_h: y1 = (border_h-1) - (max_h - 1) y2 = (border_h-1) # W diff_w = max_w - w dw1 = diff_w//2 dw2 = diff_w - dw1 x1 = xmin - dw1 x2 = xmax + dw2 if x1 < 0: x1 = 0 x2 = max_w-1 elif x2 >= border_w: x1 = (border_w-1) - (max_w - 1) x2 = (border_w-1) mask_pos_list.append( (int(x1),int(y1)) ) m = m[0][y1:y2+1,x1:x2+1] cropped_mask_list.append( m[None,...] ) return cropped_mask_list, mask_pos_list, (max_h,max_w) def crop_frames(pos_list, crop_size_hw, frame_dir): h,w = crop_size_hw for i,pos in tqdm(enumerate(pos_list),total=len(pos_list)): filename = f"{i:08d}.png" frame_path = frame_dir / filename if not frame_path.is_file(): logger.info(f"{frame_path=} not found. skip") continue if pos is None: continue x, y = pos tmp = np.asarray(Image.open(frame_path)) tmp = tmp[y:y+h,x:x+w,...] Image.fromarray(tmp).save(frame_path) def save_crop_info(mask_pos_list, crop_size_hw, frame_size_hw, save_path): import json pos_map = {} for i, pos in enumerate(mask_pos_list): if pos is not None: pos_map[str(i)]=pos info = { "frame_height" : int(frame_size_hw[0]), "frame_width" : int(frame_size_hw[1]), "height": int(crop_size_hw[0]), "width": int(crop_size_hw[1]), "pos_map" : pos_map, } with open(save_path, mode="wt", encoding="utf-8") as f: json.dump(info, f, ensure_ascii=False, indent=4) def restore_position(mask_list, crop_info): f_h = crop_info["frame_height"] f_w = crop_info["frame_width"] h = crop_info["height"] w = crop_info["width"] pos_map = crop_info["pos_map"] for i in pos_map: x,y = pos_map[i] i = int(i) m = mask_list[i] if m is None: continue m = cv2.resize( m, (w,h) ) if len(m.shape) == 2: m = m[...,None] frame = np.zeros(shape=(f_h,f_w,m.shape[2]), dtype=np.uint8) frame[y:y+h,x:x+w,...] = m mask_list[i] = frame return mask_list def load_frame_list(frame_dir, frame_array_list, crop_info): frame_list = sorted(glob.glob( os.path.join(frame_dir, "[0-9]*.png"), recursive=False)) for f in frame_list: cur = int(Path(f).stem) frame_array_list[cur] = np.asarray(Image.open(f)) if not crop_info: logger.info(f"crop_info is not exists -> skip restore") return frame_array_list for i,f in enumerate(frame_array_list): if f is None: continue frame_array_list[i] = f frame_array_list = restore_position(frame_array_list, crop_info) return frame_array_list def create_fg(mask_token, frame_dir, output_dir, output_mask_dir, masked_area_list, box_threshold=0.3, text_threshold=0.25, bg_color=(0,255,0), mask_padding=0, groundingdino_config="config/GroundingDINO/GroundingDINO_SwinB_cfg.py", groundingdino_checkpoint="data/models/GroundingDINO/groundingdino_swinb_cogcoor.pth", sam_checkpoint="data/models/SAM/sam_hq_vit_l.pth", device="cuda", ): frame_list = sorted(glob.glob( os.path.join(frame_dir, "[0-9]*.png"), recursive=False)) with torch.no_grad(): predictor = MaskPredictor( model_config_path=groundingdino_config, model_checkpoint_path=groundingdino_checkpoint, device=device, sam_checkpoint=sam_checkpoint, box_threshold=box_threshold, text_threshold=text_threshold, ) if mask_padding != 0: kernel = np.ones((abs(mask_padding),abs(mask_padding)),np.uint8) kernel2 = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (3, 3)) for i, frame in tqdm(enumerate(frame_list),total=len(frame_list), desc=f"creating mask from {mask_token=}"): frame = Path(frame) file_name = frame.name cur_frame_no = int(frame.stem) img = Image.open(frame) mask_array = predictor(img, mask_token) mask_array = mask_array[0].astype(np.uint8) * 255 if mask_padding < 0: mask_array = cv2.erode(mask_array.astype(np.uint8),kernel,iterations = 1) elif mask_padding > 0: mask_array = cv2.dilate(mask_array.astype(np.uint8),kernel,iterations = 1) mask_array = cv2.morphologyEx(mask_array.astype(np.uint8), cv2.MORPH_OPEN, kernel2) mask_array = cv2.GaussianBlur(mask_array, (7, 7), sigmaX=3, sigmaY=3, borderType=cv2.BORDER_DEFAULT) if masked_area_list[cur_frame_no] is not None: masked_area_list[cur_frame_no] = np.where(masked_area_list[cur_frame_no] > mask_array[None,...], masked_area_list[cur_frame_no], mask_array[None,...]) #masked_area_list[cur_frame_no] = masked_area_list[cur_frame_no] | mask_array[None,...] else: masked_area_list[cur_frame_no] = mask_array[None,...] if output_mask_dir: #mask_array2 = mask_array.astype(np.uint8).clip(0,1) #mask_array2 *= 255 Image.fromarray(mask_array).save( output_mask_dir / file_name ) img_array = np.asarray(img).copy() if bg_color is not None: img_array[mask_array == 0] = bg_color img = Image.fromarray(img_array) img.save( output_dir / file_name ) return masked_area_list def dilate_mask(masked_area_list, flow_mask_dilates=8, mask_dilates=5): kernel = np.ones((flow_mask_dilates,flow_mask_dilates),np.uint8) flow_masks = [ cv2.dilate(mask[0].astype(np.uint8),kernel,iterations = 1) for mask in masked_area_list ] flow_masks = [ Image.fromarray(mask * 255) for mask in flow_masks ] kernel = np.ones((mask_dilates,mask_dilates),np.uint8) dilated_masks = [ cv2.dilate(mask[0].astype(np.uint8),kernel,iterations = 1) for mask in masked_area_list ] dilated_masks = [ Image.fromarray(mask * 255) for mask in dilated_masks ] return flow_masks, dilated_masks # adapted from https://github.com/sczhou/ProPainter/blob/main/inference_propainter.py def resize_frames(frames, size=None): if size is not None: out_size = size process_size = (out_size[0]-out_size[0]%8, out_size[1]-out_size[1]%8) frames = [f.resize(process_size) for f in frames] else: out_size = frames[0].size process_size = (out_size[0]-out_size[0]%8, out_size[1]-out_size[1]%8) if not out_size == process_size: frames = [f.resize(process_size) for f in frames] return frames, process_size, out_size def get_ref_index(mid_neighbor_id, neighbor_ids, length, ref_stride=10, ref_num=-1): ref_index = [] if ref_num == -1: for i in range(0, length, ref_stride): if i not in neighbor_ids: ref_index.append(i) else: start_idx = max(0, mid_neighbor_id - ref_stride * (ref_num // 2)) end_idx = min(length, mid_neighbor_id + ref_stride * (ref_num // 2)) for i in range(start_idx, end_idx, ref_stride): if i not in neighbor_ids: if len(ref_index) > ref_num: break ref_index.append(i) return ref_index def create_bg(frame_dir, output_dir, masked_area_list, use_half = True, raft_iter = 20, subvideo_length=80, neighbor_length=10, ref_stride=10, device="cuda", low_vram = False, ): import sys repo_path = Path("src/animatediff/repo/ProPainter").absolute() repo_path = str(repo_path) sys.path.append(repo_path) from animatediff.repo.ProPainter.core.utils import to_tensors from animatediff.repo.ProPainter.model.modules.flow_comp_raft import \ RAFT_bi from animatediff.repo.ProPainter.model.propainter import InpaintGenerator from animatediff.repo.ProPainter.model.recurrent_flow_completion import \ RecurrentFlowCompleteNet from animatediff.repo.ProPainter.utils.download_util import \ load_file_from_url pretrain_model_url = 'https://github.com/sczhou/ProPainter/releases/download/v0.1.0/' model_dir = Path("data/models/ProPainter") model_dir.mkdir(parents=True, exist_ok=True) frame_list = sorted(glob.glob( os.path.join(frame_dir, "[0-9]*.png"), recursive=False)) frames = [Image.open(f) for f in frame_list] if low_vram: org_size = frames[0].size _w, _h = frames[0].size if max(_w, _h) > 512: _w = int(_w * 0.75) _h = int(_h * 0.75) frames, size, out_size = resize_frames(frames, (_w, _h)) out_size = org_size masked_area_list = [m[0] for m in masked_area_list] masked_area_list = [cv2.resize(m.astype(np.uint8), dsize=size) for m in masked_area_list] masked_area_list = [ m>127 for m in masked_area_list] masked_area_list = [m[None,...] for m in masked_area_list] else: frames, size, out_size = resize_frames(frames, None) masked_area_list = [ m>127 for m in masked_area_list] w, h = size flow_masks,masks_dilated = dilate_mask(masked_area_list) frames_inp = [np.array(f).astype(np.uint8) for f in frames] frames = to_tensors()(frames).unsqueeze(0) * 2 - 1 flow_masks = to_tensors()(flow_masks).unsqueeze(0) masks_dilated = to_tensors()(masks_dilated).unsqueeze(0) frames, flow_masks, masks_dilated = frames.to(device), flow_masks.to(device), masks_dilated.to(device) ############################################## # set up RAFT and flow competition model ############################################## ckpt_path = load_file_from_url(url=os.path.join(pretrain_model_url, 'raft-things.pth'), model_dir=model_dir, progress=True, file_name=None) fix_raft = RAFT_bi(ckpt_path, device) ckpt_path = load_file_from_url(url=os.path.join(pretrain_model_url, 'recurrent_flow_completion.pth'), model_dir=model_dir, progress=True, file_name=None) fix_flow_complete = RecurrentFlowCompleteNet(ckpt_path) for p in fix_flow_complete.parameters(): p.requires_grad = False fix_flow_complete.to(device) fix_flow_complete.eval() ############################################## # set up ProPainter model ############################################## ckpt_path = load_file_from_url(url=os.path.join(pretrain_model_url, 'ProPainter.pth'), model_dir=model_dir, progress=True, file_name=None) model = InpaintGenerator(model_path=ckpt_path).to(device) model.eval() ############################################## # ProPainter inference ############################################## video_length = frames.size(1) logger.info(f'\nProcessing: [{video_length} frames]...') with torch.no_grad(): # ---- compute flow ---- if max(w,h) <= 640: short_clip_len = 12 elif max(w,h) <= 720: short_clip_len = 8 elif max(w,h) <= 1280: short_clip_len = 4 else: short_clip_len = 2 # use fp32 for RAFT if frames.size(1) > short_clip_len: gt_flows_f_list, gt_flows_b_list = [], [] for f in range(0, video_length, short_clip_len): end_f = min(video_length, f + short_clip_len) if f == 0: flows_f, flows_b = fix_raft(frames[:,f:end_f], iters=raft_iter) else: flows_f, flows_b = fix_raft(frames[:,f-1:end_f], iters=raft_iter) gt_flows_f_list.append(flows_f) gt_flows_b_list.append(flows_b) torch.cuda.empty_cache() gt_flows_f = torch.cat(gt_flows_f_list, dim=1) gt_flows_b = torch.cat(gt_flows_b_list, dim=1) gt_flows_bi = (gt_flows_f, gt_flows_b) else: gt_flows_bi = fix_raft(frames, iters=raft_iter) torch.cuda.empty_cache() if use_half: frames, flow_masks, masks_dilated = frames.half(), flow_masks.half(), masks_dilated.half() gt_flows_bi = (gt_flows_bi[0].half(), gt_flows_bi[1].half()) fix_flow_complete = fix_flow_complete.half() model = model.half() # ---- complete flow ---- flow_length = gt_flows_bi[0].size(1) if flow_length > subvideo_length: pred_flows_f, pred_flows_b = [], [] pad_len = 5 for f in range(0, flow_length, subvideo_length): s_f = max(0, f - pad_len) e_f = min(flow_length, f + subvideo_length + pad_len) pad_len_s = max(0, f) - s_f pad_len_e = e_f - min(flow_length, f + subvideo_length) pred_flows_bi_sub, _ = fix_flow_complete.forward_bidirect_flow( (gt_flows_bi[0][:, s_f:e_f], gt_flows_bi[1][:, s_f:e_f]), flow_masks[:, s_f:e_f+1]) pred_flows_bi_sub = fix_flow_complete.combine_flow( (gt_flows_bi[0][:, s_f:e_f], gt_flows_bi[1][:, s_f:e_f]), pred_flows_bi_sub, flow_masks[:, s_f:e_f+1]) pred_flows_f.append(pred_flows_bi_sub[0][:, pad_len_s:e_f-s_f-pad_len_e]) pred_flows_b.append(pred_flows_bi_sub[1][:, pad_len_s:e_f-s_f-pad_len_e]) torch.cuda.empty_cache() pred_flows_f = torch.cat(pred_flows_f, dim=1) pred_flows_b = torch.cat(pred_flows_b, dim=1) pred_flows_bi = (pred_flows_f, pred_flows_b) else: pred_flows_bi, _ = fix_flow_complete.forward_bidirect_flow(gt_flows_bi, flow_masks) pred_flows_bi = fix_flow_complete.combine_flow(gt_flows_bi, pred_flows_bi, flow_masks) torch.cuda.empty_cache() # ---- image propagation ---- masked_frames = frames * (1 - masks_dilated) subvideo_length_img_prop = min(100, subvideo_length) # ensure a minimum of 100 frames for image propagation if video_length > subvideo_length_img_prop: updated_frames, updated_masks = [], [] pad_len = 10 for f in range(0, video_length, subvideo_length_img_prop): s_f = max(0, f - pad_len) e_f = min(video_length, f + subvideo_length_img_prop + pad_len) pad_len_s = max(0, f) - s_f pad_len_e = e_f - min(video_length, f + subvideo_length_img_prop) b, t, _, _, _ = masks_dilated[:, s_f:e_f].size() pred_flows_bi_sub = (pred_flows_bi[0][:, s_f:e_f-1], pred_flows_bi[1][:, s_f:e_f-1]) prop_imgs_sub, updated_local_masks_sub = model.img_propagation(masked_frames[:, s_f:e_f], pred_flows_bi_sub, masks_dilated[:, s_f:e_f], 'nearest') updated_frames_sub = frames[:, s_f:e_f] * (1 - masks_dilated[:, s_f:e_f]) + \ prop_imgs_sub.view(b, t, 3, h, w) * masks_dilated[:, s_f:e_f] updated_masks_sub = updated_local_masks_sub.view(b, t, 1, h, w) updated_frames.append(updated_frames_sub[:, pad_len_s:e_f-s_f-pad_len_e]) updated_masks.append(updated_masks_sub[:, pad_len_s:e_f-s_f-pad_len_e]) torch.cuda.empty_cache() updated_frames = torch.cat(updated_frames, dim=1) updated_masks = torch.cat(updated_masks, dim=1) else: b, t, _, _, _ = masks_dilated.size() prop_imgs, updated_local_masks = model.img_propagation(masked_frames, pred_flows_bi, masks_dilated, 'nearest') updated_frames = frames * (1 - masks_dilated) + prop_imgs.view(b, t, 3, h, w) * masks_dilated updated_masks = updated_local_masks.view(b, t, 1, h, w) torch.cuda.empty_cache() ori_frames = frames_inp comp_frames = [None] * video_length neighbor_stride = neighbor_length // 2 if video_length > subvideo_length: ref_num = subvideo_length // ref_stride else: ref_num = -1 # ---- feature propagation + transformer ---- for f in tqdm(range(0, video_length, neighbor_stride)): neighbor_ids = [ i for i in range(max(0, f - neighbor_stride), min(video_length, f + neighbor_stride + 1)) ] ref_ids = get_ref_index(f, neighbor_ids, video_length, ref_stride, ref_num) selected_imgs = updated_frames[:, neighbor_ids + ref_ids, :, :, :] selected_masks = masks_dilated[:, neighbor_ids + ref_ids, :, :, :] selected_update_masks = updated_masks[:, neighbor_ids + ref_ids, :, :, :] selected_pred_flows_bi = (pred_flows_bi[0][:, neighbor_ids[:-1], :, :, :], pred_flows_bi[1][:, neighbor_ids[:-1], :, :, :]) with torch.no_grad(): # 1.0 indicates mask l_t = len(neighbor_ids) # pred_img = selected_imgs # results of image propagation pred_img = model(selected_imgs, selected_pred_flows_bi, selected_masks, selected_update_masks, l_t) pred_img = pred_img.view(-1, 3, h, w) pred_img = (pred_img + 1) / 2 pred_img = pred_img.cpu().permute(0, 2, 3, 1).numpy() * 255 binary_masks = masks_dilated[0, neighbor_ids, :, :, :].cpu().permute( 0, 2, 3, 1).numpy().astype(np.uint8) for i in range(len(neighbor_ids)): idx = neighbor_ids[i] img = np.array(pred_img[i]).astype(np.uint8) * binary_masks[i] \ + ori_frames[idx] * (1 - binary_masks[i]) if comp_frames[idx] is None: comp_frames[idx] = img else: comp_frames[idx] = comp_frames[idx].astype(np.float32) * 0.5 + img.astype(np.float32) * 0.5 comp_frames[idx] = comp_frames[idx].astype(np.uint8) torch.cuda.empty_cache() # save each frame for idx in range(video_length): f = comp_frames[idx] f = cv2.resize(f, out_size, interpolation = cv2.INTER_CUBIC) f = cv2.cvtColor(f, cv2.COLOR_BGR2RGB) dst_img_path = output_dir.joinpath( f"{idx:08d}.png" ) cv2.imwrite(str(dst_img_path), f) sys.path.remove(repo_path)