hylee commited on
Commit
3d81729
·
1 Parent(s): 55c7312
APDrawingGAN2/data/single_dataset.py CHANGED
@@ -111,7 +111,7 @@ class SingleDataset(BaseDataset):
111
  cmask0 = (cmask0 >= 0.5).float()
112
  cmasks0.append(cmask0)
113
  cmask = cmask0.clone()
114
- cmask = cmask[:,center[i,1]-rhs[i]/2:center[i,1]+rhs[i]/2,center[i,0]-rws[i]/2:center[i,0]+rws[i]/2]
115
  cmasks.append(cmask)
116
  item['cmaskel'] = cmasks[0]
117
  item['cmasker'] = cmasks[1]
@@ -121,7 +121,7 @@ class SingleDataset(BaseDataset):
121
  output_nc = self.opt.output_nc
122
  mask = torch.ones([output_nc,A.shape[1],A.shape[2]])
123
  for i in range(4):
124
- mask[:,center[i,1]-rhs[i]/2:center[i,1]+rhs[i]/2,center[i,0]-rws[i]/2:center[i,0]+rws[i]/2] = 0
125
  if self.opt.soft_border:
126
  imgsize = self.opt.fineSize
127
  maskn = mask[0].numpy()
 
111
  cmask0 = (cmask0 >= 0.5).float()
112
  cmasks0.append(cmask0)
113
  cmask = cmask0.clone()
114
+ cmask = cmask[:,int(center[i,1]-rhs[i]/2):int(center[i,1]+rhs[i]/2),int(center[i,0]-rws[i]/2):int(center[i,0]+rws[i]/2)]
115
  cmasks.append(cmask)
116
  item['cmaskel'] = cmasks[0]
117
  item['cmasker'] = cmasks[1]
 
121
  output_nc = self.opt.output_nc
122
  mask = torch.ones([output_nc,A.shape[1],A.shape[2]])
123
  for i in range(4):
124
+ mask[:,int(center[i,1]-rhs[i]/2):int(center[i,1]+rhs[i]/2),int(center[i,0]-rws[i]/2):int(center[i,0]+rws[i]/2)] = 0
125
  if self.opt.soft_border:
126
  imgsize = self.opt.fineSize
127
  maskn = mask[0].numpy()