small updates

This commit is contained in:
2023-07-08 11:54:21 +02:00
parent c41a577672
commit 57b7aed4ba
15 changed files with 884 additions and 221 deletions
+4 -3
View File
@@ -66,8 +66,8 @@ class Pyramid(nn.Module):
# pred = pred + levels_input[0][2] # uv_body_posed_32
if 1 != len(self.levels):
pred = self.levels[0].enhance(pred, levels_input[0][-2]) # levels_input[i][-2]=uv_mask_res
pred = self.resizers[self.levels[1].size](pred)
out_levels.append([pred, mean, logvar])
pred = self.resizers[self.levels[1].size](pred)
pred = pred * levels_input[1][-2][:,None,:,:]
else:
pred = self.levels[0].enhance(pred, levels_input[0][-2]) # levels_input[i][-2]=uv_mask_res
@@ -86,8 +86,9 @@ class Pyramid(nn.Module):
pred = self.resizers[self.levels[i+1].size](pred)
pred = pred * levels_input[i+1][-2][:,None,:,:]
else:
pred = self.levels[-1].enhance(pred, levels_input[-1][-2])
pred = self.resizers[512](pred)
if self.levels[-1].size != 512:
pred = self.levels[-1].enhance(pred, levels_input[-1][-2])
pred = self.resizers[512](pred)
pred = pred * levels_input[-1][-1][:,None,:,:] # full res uv mask
out_levels.append([pred, mean, logvar])