Spaces:
Running
Running
File size: 2,989 Bytes
9d0b4d9 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 | import cv2
import sys
sys.path.append("..")
import util.image_processing as impro
from util import mosaic
from util import data
import torch
import numpy as np
def run_segment(img,net,size = 360,gpu_id = '-1'):
img = impro.resize(img,size)
img = data.im2tensor(img,gpu_id = gpu_id, bgr2rgb = False, is0_1 = True)
mask = net(img)
mask = data.tensor2im(mask, gray=True, is0_1 = True)
return mask
def run_pix2pix(img,net,opt):
if opt.netG == 'HD':
img = impro.resize(img,512)
else:
img = impro.resize(img,128)
img = data.im2tensor(img,gpu_id=opt.gpu_id)
img_fake = net(img)
img_fake = data.tensor2im(img_fake)
return img_fake
def traditional_cleaner(img,opt):
h,w = img.shape[:2]
img = cv2.blur(img, (opt.tr_blur,opt.tr_blur))
img = img[::opt.tr_down,::opt.tr_down,:]
img = cv2.resize(img, (w,h),interpolation=cv2.INTER_LANCZOS4)
return img
def run_styletransfer(opt, net, img):
if opt.output_size != 0:
if 'resize' in opt.preprocess and 'resize_scale_width' not in opt.preprocess:
img = impro.resize(img,opt.output_size)
elif 'resize_scale_width' in opt.preprocess:
img = cv2.resize(img, (opt.output_size,opt.output_size))
img = img[0:4*int(img.shape[0]/4),0:4*int(img.shape[1]/4),:]
if 'edges' in opt.preprocess:
if opt.canny > 100:
canny_low = opt.canny-50
canny_high = np.clip(opt.canny+50,0,255)
elif opt.canny < 50:
canny_low = np.clip(opt.canny-25,0,255)
canny_high = opt.canny+25
else:
canny_low = opt.canny-int(opt.canny/2)
canny_high = opt.canny+int(opt.canny/2)
img = cv2.Canny(img,canny_low,canny_high)
if opt.only_edges:
return img
img = data.im2tensor(img,gpu_id=opt.gpu_id,gray=True)
else:
img = data.im2tensor(img,gpu_id=opt.gpu_id)
img = net(img)
img = data.tensor2im(img)
return img
def get_ROI_position(img,net,opt,keepsize=True):
mask = run_segment(img,net,size=360,gpu_id = opt.gpu_id)
mask = impro.mask_threshold(mask,opt.mask_extend,opt.mask_threshold)
if keepsize:
mask = impro.resize_like(mask, img)
x,y,halfsize,area = impro.boundingSquare(mask, 1)
return mask,x,y,halfsize,area
def get_mosaic_position(img_origin,net_mosaic_pos,opt):
h,w = img_origin.shape[:2]
mask = run_segment(img_origin,net_mosaic_pos,size=360,gpu_id = opt.gpu_id)
# mask_1 = mask.copy()
mask = impro.mask_threshold(mask,ex_mun=int(min(h,w)/20),threshold=opt.mask_threshold)
if not opt.all_mosaic_area:
mask = impro.find_mostlikely_ROI(mask)
x,y,size,area = impro.boundingSquare(mask,Ex_mul=opt.ex_mult)
#Location fix
rat = min(h,w)/360.0
x,y,size = int(rat*x),int(rat*y),int(rat*size)
x,y = np.clip(x, 0, w),np.clip(y, 0, h)
size = np.clip(size, 0, min(w-x,h-y))
# print(x,y,size)
return x,y,size,mask |