Spaces:
Running
on
T4
Running
on
T4
| import torch | |
| from kornia.morphology import dilation, closing | |
| import requests | |
| from transformers import SamModel, SamProcessor | |
| print('Loading SAM...') | |
| device = "cuda" if torch.cuda.is_available() else "cpu" | |
| model = SamModel.from_pretrained("facebook/sam-vit-huge").to(device) | |
| processor = SamProcessor.from_pretrained("facebook/sam-vit-huge") | |
| print('DONE') | |
| def build_mask(image, faces, hairs): | |
| # 1. Segmentation | |
| input_points = faces # 2D location of the face | |
| with torch.no_grad(): | |
| inputs = processor(image, input_points=input_points, return_tensors="pt").to(device) | |
| outputs = model(**inputs) | |
| masks = processor.image_processor.post_process_masks( | |
| outputs.pred_masks.cpu(), inputs["original_sizes"].cpu(), inputs["reshaped_input_sizes"].cpu() | |
| ) | |
| scores = outputs.iou_scores | |
| input_points = hairs # 2D location of the face | |
| with torch.no_grad(): | |
| inputs = processor(image, input_points=input_points, return_tensors="pt").to(device) | |
| outputs = model(**inputs) | |
| h_masks = processor.image_processor.post_process_masks( | |
| outputs.pred_masks.cpu(), inputs["original_sizes"].cpu(), inputs["reshaped_input_sizes"].cpu() | |
| ) | |
| h_scores = outputs.iou_scores | |
| # 2. Post-processing | |
| mask=masks[0][0].all(0) | h_masks[0][0].all(0) | |
| # dilation | |
| tensor = mask[None,None,:,:] | |
| kernel = torch.ones(3, 3) | |
| mask = closing(tensor, kernel)[0,0].bool() | |
| return mask | |
| def build_mask_multi(image, faces, hairs): | |
| all_masks = [] | |
| for face,hair in zip(faces,hairs): | |
| # 1. Segmentation | |
| input_points = [face] # 2D location of the face | |
| with torch.no_grad(): | |
| inputs = processor(image, input_points=input_points, return_tensors="pt").to(device) | |
| outputs = model(**inputs) | |
| masks = processor.image_processor.post_process_masks( | |
| outputs.pred_masks.cpu(), inputs["original_sizes"].cpu(), inputs["reshaped_input_sizes"].cpu() | |
| ) | |
| scores = outputs.iou_scores | |
| input_points = [hair] # 2D location of the face | |
| with torch.no_grad(): | |
| inputs = processor(image, input_points=input_points, return_tensors="pt").to(device) | |
| outputs = model(**inputs) | |
| h_masks = processor.image_processor.post_process_masks( | |
| outputs.pred_masks.cpu(), inputs["original_sizes"].cpu(), inputs["reshaped_input_sizes"].cpu() | |
| ) | |
| h_scores = outputs.iou_scores | |
| # 2. Post-processing | |
| mask=masks[0][0].all(0) | h_masks[0][0].all(0) | |
| # dilation | |
| mask_T = mask[None,None,:,:] | |
| kernel = torch.ones(3, 3) | |
| mask = closing(mask_T, kernel)[0,0].bool() | |
| all_masks.append(mask) | |
| mask = all_masks[0] | |
| for next_mask in all_masks[1:]: | |
| mask = mask | next_mask | |
| return mask |