Update README.md

README.md

@@ -54,10 +54,11 @@ model.requires_grad_(False).cuda().eval()
 img0_path = r"_data\example_images\frame1.png"
 img2_path = r"_data\example_images\frame3.png"
 
+mean = std = [0.5]*3
 transforms = Compose([
     Resize((256, 448)),
     ToTensor(),
-    Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]),
+    Normalize(mean=mean, std=std),
 ])
 
 img0 = transforms(Image.open(img0_path).convert("RGB")).unsqueeze(0).cuda()
@@ -68,10 +69,10 @@ img1 = model.reverse_process([img0, img2], tau)
 
 plt.figure(figsize=(10, 5))
 plt.subplot(1, 3, 1)
-plt.imshow(denorm(img0, mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]).squeeze().permute(1, 2, 0).cpu().numpy())
+plt.imshow(denorm(img0, mean=mean, std=std).squeeze().permute(1, 2, 0).cpu().numpy())
 plt.subplot(1, 3, 2)
-plt.imshow(denorm(img1, mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]).squeeze().permute(1, 2, 0).cpu().numpy())
+plt.imshow(denorm(img1, mean=mean, std=std).squeeze().permute(1, 2, 0).cpu().numpy())
 plt.subplot(1, 3, 3)
-plt.imshow(denorm(img2, mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]).squeeze().permute(1, 2, 0).cpu().numpy())
+plt.imshow(denorm(img2, mean=mean, std=std).squeeze().permute(1, 2, 0).cpu().numpy())
 plt.show()
 ```