changed the visualization code when identity loss is not 1 or 0

blaxe05 · Apr 3, 2017 · de90e63 · de90e63
1 parent a510346
commit de90e63
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/models/cycle_gan_model.lua b/models/cycle_gan_model.lua
@@ -266,7 +266,7 @@ end
 
 -- returns a string that describes the display plot configuration
 function CycleGANModel:DisplayPlot(opt)
-  if opt.identity == 1 then
+  if opt.identity > 0 then
     return 'errG_A,errD_A,errRec_A,errI_A,errG_B,errD_B,errRec_B,errI_B'
   else
     return 'errG_A,errD_A,errRec_A,errG_B,errD_B,errRec_B'
@@ -302,13 +302,13 @@ function CycleGANModel:GetCurrentVisuals(opt, size)
   table.insert(visuals, {img=MakeIm3(self.real_A), label='real_A'})
   table.insert(visuals, {img=MakeIm3(self.fake_B), label='fake_B'})
   table.insert(visuals, {img=MakeIm3(self.rec_A), label='rec_A'})
-  if opt.identity == 1 then
+  if opt.identity > 0 then
     table.insert(visuals, {img=MakeIm3(self.identity_A), label='identity_A'})
   end
   table.insert(visuals, {img=MakeIm3(self.real_B), label='real_B'})
   table.insert(visuals, {img=MakeIm3(self.fake_A), label='fake_A'})
   table.insert(visuals, {img=MakeIm3(self.rec_B), label='rec_B'})
-  if opt.identity == 1 then
+  if opt.identity > 0 then
     table.insert(visuals, {img=MakeIm3(self.identity_B), label='identity_B'})
   end
   return visuals