Update app.py

app.py

@@ -13,11 +13,20 @@ to_img=T.ToPILImage()
 # examples=["examples/input_0.png","examples/input_9.png"]
 example_list = [["examples/" + example] for example in os.listdir("examples")]
 # example_list=['1.jpg','2.jpg']
-def de_norm(img):
-    img_ = img.mul(torch.FloatTensor(STD).view(3, 1, 1))
-    img_ = img_.add(torch.FloatTensor(MEAN).view(3, 1, 1)).detach().numpy()
-    img_ = np.transpose(img_, (1, 2, 0))
-    return img_ 
+# def de_norm(img):
+#     img_ = img.mul(torch.FloatTensor(STD).view(3, 1, 1))
+#     img_ = img_.add(torch.FloatTensor(MEAN).view(3, 1, 1)).detach().numpy()
+#     img_ = np.transpose(img_, (1, 2, 0))
+#     return img_ 
+inverse_transform = transforms.Compose([    transforms.Normalize(mean=[-0.5, -0.5, -0.5], std=[1/0.5, 1/0.5, 1/0.5]),
+    transforms.Normalize(mean=[-transform.mean[0]/transform.std[0],
+                               -transform.mean[1]/transform.std[1],
+                               -transform.mean[2]/transform.std[2]],
+                          std=[1/transform.std[0],
+                               1/transform.std[1],
+                               1/transform.std[2]]),
+    transforms.ToPILImage()
+])
 def predict(img):
     # Apply Transformations
     img = transform_gen(img).unsqueeze(0)
@@ -27,8 +36,7 @@ def predict(img):
     with torch.inference_mode():
         y_gen = gen(img)
     y_gen = y_gen[0]
-    y_gen=de_norm(y_gen)
-    y_gen = to_img(y_gen)
+    y_gen = inverse_transform(y_gen)
 
     return y_gen