Internal change.

PiperOrigin-RevId: 317682996
google · Jun 22, 2020 · 2e2f38c · 2e2f38c
1 parent 65b8739
commit 2e2f38c
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 2 deletions.
diff --git a/nucleus/util/BUILD b/nucleus/util/BUILD
@@ -370,7 +370,10 @@ cc_library(
 py_library(
     name = "vis",
     srcs = ["vis.py"],
-    deps = ["//nucleus/protos:variants_py_pb2"],
+    deps = [
+        "//nucleus/io:gfile",
+        "//nucleus/protos:variants_py_pb2",
+    ],
 )
 
 py_test(

diff --git a/nucleus/util/vis.py b/nucleus/util/vis.py
@@ -23,6 +23,7 @@
 from __future__ import division
 from __future__ import print_function
 
+from nucleus.io import gfile
 from IPython import display
 import numpy as np
 from PIL import Image
@@ -297,7 +298,8 @@ def save_to_png(arr,
   # Saving to a temporary file is needed even when showing in a notebook
   if path is None:
     path = '/tmp/tmp.png'
-  img.save(path)
+  with gfile.Open(path, 'wb') as fout:
+    img.save(fout, format=path.split('.')[-1])
 
   # Show image (great for notebooks)
   if show: