Spaces:
Running
Running
| import cv2 | |
| import numpy as np | |
| import torch | |
| import threading | |
| from torchvision import transforms | |
| from clip.clipseg import CLIPDensePredT | |
| import numpy as np | |
| from roop.typing import Frame | |
| THREAD_LOCK_CLIP = threading.Lock() | |
| class Mask_Clip2Seg(): | |
| plugin_options:dict = None | |
| model_clip = None | |
| processorname = 'clip2seg' | |
| type = 'mask' | |
| def Initialize(self, plugin_options:dict): | |
| if self.plugin_options is not None: | |
| if self.plugin_options["devicename"] != plugin_options["devicename"]: | |
| self.Release() | |
| self.plugin_options = plugin_options | |
| if self.model_clip is None: | |
| self.model_clip = CLIPDensePredT(version='ViT-B/16', reduce_dim=64, complex_trans_conv=True) | |
| self.model_clip.eval(); | |
| self.model_clip.load_state_dict(torch.load('models/CLIP/rd64-uni-refined.pth', map_location=torch.device('cpu')), strict=False) | |
| device = torch.device(self.plugin_options["devicename"]) | |
| self.model_clip.to(device) | |
| def Run(self, img1, keywords:str) -> Frame: | |
| if keywords is None or len(keywords) < 1 or img1 is None: | |
| return img1 | |
| source_image_small = cv2.resize(img1, (256,256)) | |
| img_mask = np.full((source_image_small.shape[0],source_image_small.shape[1]), 0, dtype=np.float32) | |
| mask_border = 1 | |
| l = 0 | |
| t = 0 | |
| r = 1 | |
| b = 1 | |
| mask_blur = 5 | |
| clip_blur = 5 | |
| img_mask = cv2.rectangle(img_mask, (mask_border+int(l), mask_border+int(t)), | |
| (256 - mask_border-int(r), 256-mask_border-int(b)), (255, 255, 255), -1) | |
| img_mask = cv2.GaussianBlur(img_mask, (mask_blur*2+1,mask_blur*2+1), 0) | |
| img_mask /= 255 | |
| input_image = source_image_small | |
| transform = transforms.Compose([ | |
| transforms.ToTensor(), | |
| transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), | |
| transforms.Resize((256, 256)), | |
| ]) | |
| img = transform(input_image).unsqueeze(0) | |
| thresh = 0.5 | |
| prompts = keywords.split(',') | |
| with THREAD_LOCK_CLIP: | |
| with torch.no_grad(): | |
| preds = self.model_clip(img.repeat(len(prompts),1,1,1), prompts)[0] | |
| clip_mask = torch.sigmoid(preds[0][0]) | |
| for i in range(len(prompts)-1): | |
| clip_mask += torch.sigmoid(preds[i+1][0]) | |
| clip_mask = clip_mask.data.cpu().numpy() | |
| np.clip(clip_mask, 0, 1) | |
| clip_mask[clip_mask>thresh] = 1.0 | |
| clip_mask[clip_mask<=thresh] = 0.0 | |
| kernel = np.ones((5, 5), np.float32) | |
| clip_mask = cv2.dilate(clip_mask, kernel, iterations=1) | |
| clip_mask = cv2.GaussianBlur(clip_mask, (clip_blur*2+1,clip_blur*2+1), 0) | |
| img_mask *= clip_mask | |
| img_mask[img_mask<0.0] = 0.0 | |
| return img_mask | |
| def Release(self): | |
| self.model_clip = None | |