From 15d7ba077ed0817fb505e35a0b69bff3a478b81e Mon Sep 17 00:00:00 2001
From: Rayce Rossum <raycer>
Date: Sun, 17 Feb 2019 14:42:13 -0800
Subject: [PATCH] Update translate for consistency w/ other methods

---
 AutoTransformPy/translate.py | 6 +++---
 tests/test_translate.py      | 6 +++---
 2 files changed, 6 insertions(+), 6 deletions(-)

diff --git a/AutoTransformPy/translate.py b/AutoTransformPy/translate.py
index 93e0fa1..6582797 100644
--- a/AutoTransformPy/translate.py
+++ b/AutoTransformPy/translate.py
@@ -2,7 +2,7 @@
 from skimage.transform import AffineTransform, warp
 import numpy as np
 
-from matplotlib import pyplot as plt
+# from matplotlib import pyplot as plt
 
 def translate (image_path, num_images, max_translation):
     """Returns an array of images of length num_images randomly translated a random number of pixels up to max_rotation
@@ -47,13 +47,13 @@ def translate (image_path, num_images, max_translation):
 
     # Perform image translation
     translations = (np.random.randint(-max_translation, max_translation, num_images), np.random.randint(-max_translation, max_translation, num_images))
-    translated_images = []
+    translated_images = [org_image]
 
     for i in range(len(translations[0])):
         tform = AffineTransform(translation=(translations[0][i], translations[1][i]))
         translated_images.append(warp(org_image, tform, mode="constant", preserve_range=True).astype('uint8'))
 
-    return np.asarray(translated_images), org_image
+    return np.asarray(translated_images)
 
 
 # data, org_image = translate("../tests/imgs/milad.jpg", 10, 50)
diff --git a/tests/test_translate.py b/tests/test_translate.py
index 5f109ae..d88aac6 100644
--- a/tests/test_translate.py
+++ b/tests/test_translate.py
@@ -15,8 +15,8 @@ def test_inputs():
 
 def test_return_imgs(): # Tests that the number of images returned from translate is correct
     test_img = imread("../tests/imgs/milad.jpg")
-    returned_arr, org_image = trans.translate("../tests/imgs/milad.jpg", 5, 10)
+    returned_arr = trans.translate("../tests/imgs/milad.jpg", 5, 10)
 
-    assert org_image.shape == test_img.shape
-    assert returned_arr.shape[0] == 5
+    assert returned_arr[0].shape == test_img.shape
+    assert returned_arr.shape[0] == 6
     assert returned_arr.shape[1:] == test_img.shape