Update README.md

README.md

@@ -95,6 +95,43 @@ matte = (matte * 255).astype(np.uint8)
 Image.fromarray(matte).save("alpha_matte.png")
 
 ```
+## Visualization
+
+```python
+import numpy as np
+from PIL import Image
+
+def combined_display(image, matte):
+  # calculate display resolution
+  w, h = image.width, image.height
+  rw, rh = 800, int(h * 800 / (3 * w))
+  
+  # obtain predicted foreground
+  image = np.asarray(image)
+  if len(image.shape) == 2:
+    image = image[:, :, None]
+  if image.shape[2] == 1:
+    image = np.repeat(image, 3, axis=2)
+  elif image.shape[2] == 4:
+    image = image[:, :, 0:3]
+  matte = np.repeat(np.asarray(matte)[:, :, None], 3, axis=2) / 255
+  foreground = image * matte + np.full(image.shape, 255) * (1 - matte)
+  
+  # combine image, foreground, and alpha into one line
+  combined = np.concatenate((image, foreground, matte * 255), axis=1)
+  combined = Image.fromarray(np.uint8(combined)).resize((rw, rh))
+  return combined
+
+# visualize all images
+image_names = os.listdir(input_folder)
+for image_name in image_names:
+  matte_name = image_name.split('.')[0] + '.png'
+  image = Image.open(os.path.join(input_folder, image_name))
+  matte = Image.open(os.path.join(output_folder, matte_name))
+  display(combined_display(image, matte))
+  print(image_name, '\n')
+
+```
 
 
 ---