Spaces:

gblinc111
/

Intelligent-Photo-Blur-Using-Dichotomous-Image-Segmentation

Runtime error

App Files Files Community

Intelligent-Photo-Blur-Using-Dichotomous-Image-Segmentation / app.py

gblinc111

Fix DIS hyperlink (#1)

82bc132 over 1 year ago

raw history blame contribute delete

No virus

7.09 kB

	import gradio as gr
	from PIL import Image, ImageFilter
	import numpy as np
	import torch
	from torch.autograd import Variable
	from torchvision import transforms
	import torch.nn.functional as F
	import gdown
	import os

	os.system("git clone https://github.com/xuebinqin/DIS")
	os.system("mv DIS/IS-Net/* .")

	# project imports
	from data_loader_cache import normalize, im_reader, im_preprocess
	from models import *

	#Helpers
	device = 'cuda' if torch.cuda.is_available() else 'cpu'

	# Download official weights
	if not os.path.exists("saved_models"):
	os.mkdir("saved_models")
	MODEL_PATH_URL = "https://drive.google.com/uc?id=1KyMpRjewZdyYfxHPYcd-ZbanIXtin0Sn"
	gdown.download(MODEL_PATH_URL, "saved_models/isnet.pth", use_cookies=False)

	class GOSNormalize(object):
	'''
	Normalize the Image using torch.transforms
	'''
	def __init__(self, mean=[0.485,0.456,0.406], std=[0.229,0.224,0.225]):
	self.mean = mean
	self.std = std

	def __call__(self,image):
	image = normalize(image,self.mean,self.std)
	return image


	transform = transforms.Compose([GOSNormalize([0.5,0.5,0.5],[1.0,1.0,1.0])])

	def load_image(im_path, hypar):
	im = im_reader(im_path)
	im, im_shp = im_preprocess(im, hypar["cache_size"])
	im = torch.divide(im,255.0)
	shape = torch.from_numpy(np.array(im_shp))
	return transform(im).unsqueeze(0), shape.unsqueeze(0) # make a batch of image, shape


	def build_model(hypar,device):
	net = hypar["model"]#GOSNETINC(3,1)

	# convert to half precision
	if(hypar["model_digit"]=="half"):
	net.half()
	for layer in net.modules():
	if isinstance(layer, nn.BatchNorm2d):
	layer.float()

	net.to(device)

	if(hypar["restore_model"]!=""):
	net.load_state_dict(torch.load(hypar["model_path"]+"/"+hypar["restore_model"], map_location=device))
	net.to(device)
	net.eval()
	return net


	def predict(net, inputs_val, shapes_val, hypar, device):
	'''
	Given an Image, predict the mask
	'''
	net.eval()

	if(hypar["model_digit"]=="full"):
	inputs_val = inputs_val.type(torch.FloatTensor)
	else:
	inputs_val = inputs_val.type(torch.HalfTensor)


	inputs_val_v = Variable(inputs_val, requires_grad=False).to(device) # wrap inputs in Variable

	ds_val = net(inputs_val_v)[0] # list of 6 results

	pred_val = ds_val[0][0,:,:,:] # B x 1 x H x W # we want the first one which is the most accurate prediction

	## recover the prediction spatial size to the orignal image size
	pred_val = torch.squeeze(F.upsample(torch.unsqueeze(pred_val,0),(shapes_val[0][0],shapes_val[0][1]),mode='bilinear'))

	ma = torch.max(pred_val)
	mi = torch.min(pred_val)
	pred_val = (pred_val-mi)/(ma-mi) # max = 1

	if device == 'cuda': torch.cuda.empty_cache()
	return (pred_val.detach().cpu().numpy()*255).astype(np.uint8) # it is the mask we need

	# Set Parameters
	hypar = {} # paramters for inferencing


	hypar["model_path"] ="./saved_models" ## load trained weights from this path
	hypar["restore_model"] = "isnet.pth" ## name of the to-be-loaded weights
	hypar["interm_sup"] = False ## indicate if activate intermediate feature supervision

	## choose floating point accuracy --
	hypar["model_digit"] = "full" ## indicates "half" or "full" accuracy of float number
	hypar["seed"] = 0

	hypar["cache_size"] = [1024, 1024] ## cached input spatial resolution, can be configured into different size

	## data augmentation parameters ---
	hypar["input_size"] = [1024, 1024] ## mdoel input spatial size, usually use the same value hypar["cache_size"], which means we don't further resize the images
	hypar["crop_size"] = [1024, 1024] ## random crop size from the input, it is usually set as smaller than hypar["cache_size"], e.g., [920,920] for data augmentation

	hypar["model"] = ISNetDIS()

	# Build Model
	net = build_model(hypar, device)


	def infer_mask(image: Image):
	image_path = image

	image_tensor, orig_size = load_image(image_path, hypar)
	mask = predict(net, image_tensor, orig_size, hypar, device)

	return Image.fromarray(mask).convert("L")

	def blur(image_set: list, blur_amount: int):
	blurred_image = image_set[0].filter(ImageFilter.GaussianBlur(blur_amount))

	return Image.composite(image_set[0], blurred_image, image_set[1])


	with gr.Blocks() as interface:
	default_im = Image.open("newman.jpg").convert("RGB")
	default_mask = Image.open("newman_mask.jpg").convert("RGB")
	examples_list = [os.path.join(os.path.dirname(__file__), "newman.jpg"),
	os.path.join(os.path.dirname(__file__), "abbey.jpg"),
	os.path.join(os.path.dirname(__file__), "julia.jpg")
	]

	current_images = gr.State([default_im, default_mask])
	mask_toggle = gr.State(False)

	gr.Markdown(
	"""
	### Intelligent Photo Blur Using Dichotomous Image Segmentation

	This app leverages the machine learning engine built by Xuebin Qin ([https://github.com/xuebinqin/DIS](https://github.com/xuebinqin/DIS)) to mask the prominent subject within a photograph.
	The mask is used to keep the subject in clear focus while an adjustable slider is available to interactively blur the background.
	To use, upload a photo and press the run button. You can adjust the level of blur through the slider and view the mask using the "Show Generated Mask" button.
	"""
	)
	with gr.Row():
	with gr.Column():
	input_image = gr.Image(value=default_im, type='filepath')
	run_button = gr.Button()
	gr.Examples(inputs=input_image, examples=examples_list)
	with gr.Column():
	output_image = gr.Image()
	blur_slider = gr.Slider(0, 16, 5, step=1, label="Blur Amount")
	mask_button = gr.Button(value="Show Generated Mask")
	mask_image = gr.Image(value=default_mask, visible=False)

	def run(image: Image, current_images: gr.State):
	im_rgb = Image.open(image).convert("RGB")
	mask = infer_mask(image)

	return (
	blur([im_rgb, mask], 5),
	mask,
	[im_rgb, mask]
	)

	def reset_slider():
	return gr.update(value=5)

	def show_mask(mask_toggle: gr.State):
	if mask_toggle == True:
	return gr.update(visible=False)
	else:
	return gr.update(visible=True)

	def toggle_mask(mask_toggle: gr.State):
	if mask_toggle == True:
	return False
	else:
	return True

	run_button.click(run, [input_image, current_images], [output_image, mask_image, current_images])
	run_button.click(reset_slider, outputs=blur_slider)
	blur_slider.change(blur, [current_images, blur_slider], output_image, show_progress=False)
	mask_button.click(show_mask, inputs=mask_toggle, outputs=mask_image)
	mask_button.click(toggle_mask, inputs=mask_toggle, outputs=mask_toggle)

	interface.launch()