第Python照片人物背景替換的實(shí)現(xiàn)方法#ForTorchScriptexportoptimization.

__constants__=['kernel_size','patch_crop_method','patch_replace_method']

def__init__(self,

mode:str,

sample_pixels:int,

threshold:float,

kernel_size:int=3,

prevent_oversampling:bool=True,

patch_crop_method:str='unfold',

patch_replace_method:str='scatter_nd'):

super().__init__()

assertmodein['full','sampling','thresholding']

assertkernel_sizein[1,3]

assertpatch_crop_methodin['unfold','roi_align','gather']

assertpatch_replace_methodin['scatter_nd','scatter_element']

self.mode=mode

self.sample_pixels=sample_pixels

self.threshold=threshold

self.kernel_size=kernel_size

self.prevent_oversampling=prevent_oversampling

self.patch_crop_method=patch_crop_method

self.patch_replace_method=patch_replace_method

channels=[32,24,16,12,4]

self.conv1=nn.Conv2d(channels[0]+6+4,channels[1],kernel_size,bias=False)

self.bn1=nn.BatchNorm2d(channels[1])

self.conv2=nn.Conv2d(channels[1],channels[2],kernel_size,bias=False)

self.bn2=nn.BatchNorm2d(channels[2])

self.conv3=nn.Conv2d(channels[2]+6,channels[3],kernel_size,bias=False)

self.bn3=nn.BatchNorm2d(channels[3])

self.conv4=nn.Conv2d(channels[3],channels[4],kernel_size,bias=True)

self.relu=nn.ReLU(True)

defforward(self,

src:torch.Tensor,

bgr:torch.Tensor,

pha:torch.Tensor,

fgr:torch.Tensor,

err:torch.Tensor,

hid:torch.Tensor):

H_full,W_full=src.shape[2:]

H_half,W_half=H_full//2,W_full//2

H_quat,W_quat=H_full//4,W_full//4

src_bgr=torch.cat([src,bgr],dim=1)

ifself.mode!='full':

err=F.interpolate(err,(H_quat,W_quat),mode='bilinear',align_corners=False)

ref=self.select_refinement_regions(err)

idx=torch.nonzero(ref.squeeze(1))

idx=idx[:,0],idx[:,1],idx[:,2]

ifidx[0].size(0)0:

x=torch.cat([hid,pha,fgr],dim=1)

x=F.interpolate(x,(H_half,W_half),mode='bilinear',align_corners=False)

x=self.crop_patch(x,idx,2,3ifself.kernel_size==3else0)

y=F.interpolate(src_bgr,(H_half,W_half),mode='bilinear',align_corners=False)

y=self.crop_patch(y,idx,2,3ifself.kernel_size==3else0)

x=self.conv1(torch.cat([x,y],dim=1))

x=self.bn1(x)

x=self.relu(x)

x=self.conv2(x)

x=self.bn2(x)

x=self.relu(x)

x=F.interpolate(x,8ifself.kernel_size==3else4,mode='nearest')

y=self.crop_patch(src_bgr,idx,4,2ifself.kernel_size==3else0)

x=self.conv3(torch.cat([x,y],dim=1))

x=self.bn3(x)

x=self.relu(x)

x=self.conv4(x)

out=torch.cat([pha,fgr],dim=1)

out=F.interpolate(out,(H_full,W_full),mode='bilinear',align_corners=False)

out=self.replace_patch(out,x,idx)

pha=out[:,:1]

fgr=out[:,1:]

else:

pha=F.interpolate(pha,(H_full,W_full),mode='bilinear',align_corners=False)

fgr=F.interpolate(fgr,(H_full,W_full),mode='bilinear',align_corners=False)

else:

x=torch.cat([hid,pha,fgr],dim=1)

x=F.interpolate(x,(H_half,W_half),mode='bilinear',align_corners=False)

y=F.interpolate(src_bgr,(H_half,W_half),mode='bilinear',align_corners=False)

ifself.kernel_size==3:

x=F.pad(x,(3,3,3,3))

y=F.pad(y,(3,3,3,3))

x=self.conv1(torch.cat([x,y],dim=1))

x=self.bn1(x)

x=self.relu(x)

x=self.conv2(x)

x=self.bn2(x)

x=self.relu(x)

ifself.kernel_size==3:

x=F.interpolate(x,(H_full+4,W_full+4))

y=F.pad(src_bgr,(2,2,2,2))

else:

x=F.interpolate(x,(H_full,W_full),mode='nearest')

y=src_bgr

x=self.conv3(torch.cat([x,y],dim=1))

x=self.bn3(x)

x=self.relu(x)

x=self.conv4(x)

pha=x[:,:1]

fgr=x[:,1:]

ref=torch.ones((src.size(0),1,H_quat,W_quat),device=src.device,dtype=src.dtype)

returnpha,fgr,ref

defselect_refinement_regions(self,err:torch.Tensor):

Selectrefinementregions.

Input:

err:errormap(B,1,H,W)

Output:

ref:refinementregions(B,1,H,W).FloatTensor.1isselected,0isnot.

ifself.mode=='sampling':

#Samplingmode.

b,_,h,w=err.shape

err=err.view(b,-1)

idx=err.topk(self.sample_pixels//16,dim=1,sorted=False).indices

ref=torch.zeros_like(err)

ref.scatter_(1,idx,1.)

ifself.prevent_oversampling:

ref.mul_(err.gt(0).float())

ref=ref.view(b,1,h,w)

else:

#Thresholdingmode.

ref=err.gt(self.threshold).float()

returnref

defcrop_patch(self,

x:torch.Tensor,

idx:Tuple[torch.Tensor,torch.Tensor,torch.Tensor],

size:int,

padding:int):

Cropsselectedpatchesfromimagegivenindices.

Inputs:

x:image(B,C,H,W).

idx:selectionindicesTuple[(P,),(P,),(P,),],wherethe3valuesare(B,H,W)index.

size:centersizeofthepatch,alsostrideofthecrop.

padding:expansionsizeofthepatch.

Output:

patch:(P,C,h,w),whereh=w=size+2*padding.

ifpadding!=0:

x=F.pad(x,(padding,)*4)

ifself.patch_crop_method=='unfold':

#Useunfold.BestperformanceforPyTorchandTorchScript.

returnx.permute(0,2,3,1)\

.unfold(1,size+2*padding,size)\

.unfold(2,size+2*padding,size)[idx[0],idx[1],idx[2]]

elifself.patch_crop_method=='roi_align':

#Useroi_align.BestcompatibilityforONNX.

idx=idx[0].type_as(x),idx[1].type_as(x),idx[2].type_as(x)

b=idx[0]

x1=idx[2]*size-0.5

y1=idx[1]*size-0.5

x2=idx[2]*size+size+2*padding-0.5

y2=idx[1]*size+size+2*padding-0.5

boxes=torch.stack([b,x1,y1,x2,y2],dim=1)

returntorchvision.ops.roi_align(x,boxes,size+2*padding,sampling_ratio=1)

else:

#Usegather.Cropsoutpatchespixelbypixel.

idx_pix=pute_pixel_indices(x,idx,size,padding)

pat=torch.gather(x.view(-1),0,idx_pix.view(-1))

pat=pat.view(-1,x.size(1),size+2*padding,size+2*padding)

returnpat

defreplace_patch(self,

x:torch.Tensor,

y:torch.Tensor,

idx:Tuple[torch.Tensor,torch.Tensor,torch.Tensor]):

Replacespatchesbackintoimagegivenindex.

Inputs:

x:image(B,C,H,W)

y:patches(P,C,h,w)

idx:selectionindicesTuple[(P,),(P,),(P,)]wherethe3valuesare(B,H,W)index.

Output:

image:(B,C,H,W),wherepatchesatidxlocationsarereplacedwithy.

xB,xC,xH,xW=x.shape

yB,yC,yH,yW=y.shape

ifself.patch_replace_method=='scatter_nd':

#Usescatter_nd.BestperformanceforPyTorchandTorchScript.Replacingpatchbypatch.

x=x.view(xB,xC,xH//yH,yH,xW//yW,yW).permute(0,2,4,1,3,5)

x[idx[0],idx[1],idx[2]]=y

x=x.permute(0,3,1,4,2,5).view(xB,xC,xH,xW)

returnx

else:

#Usescatter_element.BestcompatibilityforONNX.Replacingpixelbypixel.

idx_pix=pute_pixel_indices(x,idx,size=4,padding=0)

returnx.view(-1).scatter_(0,idx_pix.view(-1),y.view(-1)).view(x.shape)

defcompute_pixel_indices(self,

x:torch.Tensor,

idx:Tuple[torch.Tensor,torch.Tensor,torch.Tensor],

size:int,

padding:int):

Computeselectedpixelindicesinthetensor.

Usedforcrop_method=='gather'andreplace_method=='scatter_element',whichcropandreplacepixelbypixel.

Input:

x:image:(B,C,