1 files changed, 12 insertions, 0 deletions
diff --git a/megapixels/app/utils/im_utils.py b/megapixels/app/utils/im_utils.py
index e882c67f..d36c1c32 100644
--- a/megapixels/app/utils/im_utils.py
+++ b/megapixels/app/utils/im_utils.py
@@ -19,7 +19,19 @@ from torch.autograd import Variable
 from sklearn.metrics.pairwise import cosine_similarity
 import datetime
 
+def ensure_pil(im):
+    """Ensure image is Pillow format"""
+    try:
+        im.verify()
+        return im
+    except:
+        return Image.fromarray(im.astype('uint8'), 'RGB')
 
+def ensure_np(im):
+    """Ensure image is numpy array"""
+    if type(im) == np.ndarray:
+        return im
+    return np.asarray(im, np.uint8)
 
 def num_channels(im):
   '''Returns number of channels in numpy.ndarray image'''