fix: out mask images can be computed with more than 2 semantic classes

jolibrain · Feb 18, 2021 · 36d54b9 · 36d54b9
1 parent 233634b
commit 36d54b9
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 4 deletions.
diff --git a/models/cycle_gan_semantic_mask_model.py b/models/cycle_gan_semantic_mask_model.py
@@ -249,7 +249,7 @@ def forward(self):
             if hasattr(self,'criterionMask'):
                 label_A = self.input_A_label
                 label_A_in = label_A.unsqueeze(1)
-                label_A_inv = torch.tensor(np.ones(label_A.size())).to(self.device) - label_A
+                label_A_inv = torch.tensor(np.ones(label_A.size())).to(self.device) - label_A>0
                 label_A_inv = label_A_inv.unsqueeze(1)
                 #label_A_inv = torch.cat ([label_A_inv,label_A_inv,label_A_inv],1)
 
@@ -270,7 +270,7 @@ def forward(self):
 
                     label_B = self.input_B_label
                     label_B_in = label_B.unsqueeze(1)
-                    label_B_inv = torch.tensor(np.ones(label_B.size())).to(self.device) - label_B
+                    label_B_inv = torch.tensor(np.ones(label_B.size())).to(self.device) - label_B>0
                     label_B_inv = label_B_inv.unsqueeze(1)
 
                     self.real_B_out_mask = self.real_B *label_B_inv

diff --git a/models/cycle_gan_semantic_mask_sty2_model.py b/models/cycle_gan_semantic_mask_sty2_model.py
@@ -374,7 +374,7 @@ def forward(self):
             if hasattr(self,'criterionMask'):
                 label_A = self.input_A_label
                 label_A_in = label_A.unsqueeze(1)
-                label_A_inv = torch.tensor(np.ones(label_A.size())).to(self.device) - label_A
+                label_A_inv = torch.tensor(np.ones(label_A.size())).to(self.device) - label_A>0
                 label_A_inv = label_A_inv.unsqueeze(1)
                 #label_A_inv = torch.cat ([label_A_inv,label_A_inv,label_A_inv],1)
 
@@ -393,7 +393,7 @@ def forward(self):
 
                     label_B = self.input_B_label
                     label_B_in = label_B.unsqueeze(1)
-                    label_B_inv = torch.tensor(np.ones(label_B.size())).to(self.device) - label_B
+                    label_B_inv = torch.tensor(np.ones(label_B.size())).to(self.device) - label_B>0
                     label_B_inv = label_B_inv.unsqueeze(1)
                     #label_B_inv = torch.cat ([label_B_inv,label_B_inv,label_B_inv],1)