Spaces:

facebook
/

MaskCut

Build error

App Files Files Community

MaskCut / model.py

hysts HF staff

Add files

73cbad1 over 1 year ago

raw history blame contribute delete

No virus

3.19 kB

	# This file is adapted from https://github.com/facebookresearch/CutLER/blob/077938c626341723050a1971107af552a6ca6697/maskcut/demo.py
	# The original license file is the file named LICENSE.CutLER in this repo.

	import sys

	import numpy as np
	import PIL.Image as Image
	import torch
	from scipy import ndimage

	sys.path.append('CutLER/maskcut/')
	sys.path.append('CutLER/')
	import dino
	from colormap import random_color
	from crf import densecrf
	from maskcut import maskcut
	from third_party.TokenCut.unsupervised_saliency_detection import metric


	def vis_mask(input, mask, mask_color):
	fg = mask > 0.5
	rgb = np.copy(input)
	rgb[fg] = (rgb[fg] * 0.3 + np.array(mask_color) * 0.7).astype(np.uint8)
	return Image.fromarray(rgb)


	class Model:
	def __init__(self):
	self.device = torch.device(
	'cuda:0' if torch.cuda.is_available() else 'cpu')
	self.backbone = self.load_backbone()

	def load_backbone(self):
	# DINO hyperparameters
	vit_arch = 'base'
	vit_feat = 'k'
	patch_size = 8
	# DINO pre-trained model
	url = 'https://dl.fbaipublicfiles.com/dino/dino_vitbase8_pretrain/dino_vitbase8_pretrain.pth'
	feat_dim = 768

	# extract patch features with a pretrained DINO model
	backbone = dino.ViTFeat(url, feat_dim, vit_arch, vit_feat, patch_size)
	backbone.eval()
	backbone.to(self.device)
	return backbone

	def __call__(self, img_path, tau, n, fixed_size=480):
	# get pseudo-masks with MaskCut
	bipartitions, _, I_new = maskcut(img_path,
	self.backbone,
	self.backbone.patch_size,
	tau,
	N=n,
	fixed_size=fixed_size,
	cpu=self.device.type == 'cpu')
	I = Image.open(img_path).convert('RGB')
	width, height = I.size
	pseudo_mask_list = []
	for idx, bipartition in enumerate(bipartitions):
	# post-process pseudo-masks with CRF
	pseudo_mask = densecrf(np.array(I_new), bipartition)
	pseudo_mask = ndimage.binary_fill_holes(pseudo_mask >= 0.5)

	# filter out the mask that have a very different pseudo-mask after the CRF
	mask1 = torch.from_numpy(bipartition).to(self.device)
	mask2 = torch.from_numpy(pseudo_mask).to(self.device)
	if metric.IoU(mask1, mask2) < 0.5:
	pseudo_mask = pseudo_mask * -1

	# construct binary pseudo-masks
	pseudo_mask[pseudo_mask < 0] = 0
	pseudo_mask = Image.fromarray(np.uint8(pseudo_mask * 255))
	pseudo_mask = np.asarray(pseudo_mask.resize((width, height)))

	pseudo_mask = pseudo_mask.astype(np.uint8)
	upper = np.max(pseudo_mask)
	lower = np.min(pseudo_mask)
	thresh = upper / 2.0
	pseudo_mask[pseudo_mask > thresh] = upper
	pseudo_mask[pseudo_mask <= thresh] = lower
	pseudo_mask_list.append(pseudo_mask)
	return pseudo_mask_list