diff --git a/modern_colorthief.pyi b/modern_colorthief.pyi
index b6ec836..450de32 100644
--- a/modern_colorthief.pyi
+++ b/modern_colorthief.pyi
@@ -1,11 +1,11 @@
 __version__: str
 
 def get_color(
-    location: str,
+    image: str,
     quality: int | None = 10,
 ) -> tuple[int, int, int]: ...
 def get_palette(
-    location: str,
+    image: str,
     color_count: int | None = 10,
     quality: int | None = 10,
 ) -> list[tuple[int, int, int]]: ...
diff --git a/src/python/__init__.py b/src/python/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/src/lib.rs b/src/rust/lib.rs
similarity index 82%
rename from src/lib.rs
rename to src/rust/lib.rs
index a0ea853..99628ac 100644
--- a/src/lib.rs
+++ b/src/rust/lib.rs
@@ -14,11 +14,11 @@ fn get_image_buffer(img: image::DynamicImage) -> (Vec<u8>, ColorFormat) {
 /// Returns the pallette given an image
 #[pyfunction]
 fn get_palette(
-    location: String,
+    image: String,
     color_count: Option<u8>,
     quality: Option<u8>,
 ) -> PyResult<Vec<(u8, u8, u8)>> {
-    let img = image::open(&std::path::Path::new(&location)).unwrap();
+    let img = image::open(&std::path::Path::new(&image)).unwrap();
     let (buffer, color_type) = get_image_buffer(img);
 
     let colors = color_thief::get_palette(
@@ -39,8 +39,8 @@ fn get_palette(
 
 // Gets the dominant color given an image
 #[pyfunction]
-fn get_color(location: String, quality: Option<u8>) -> PyResult<(u8, u8, u8)> {
-    let palette = get_palette(location, Some(5), Some(quality.unwrap_or(10))).unwrap();
+fn get_color(image: String, quality: Option<u8>) -> PyResult<(u8, u8, u8)> {
+    let palette = get_palette(image, Some(5), Some(quality.unwrap_or(10))).unwrap();
     Ok(palette[0])
 }
 
diff --git a/test.jpg b/test.jpg
new file mode 100644
index 0000000..564bd05
Binary files /dev/null and b/test.jpg differ
diff --git a/test.py b/test.py
new file mode 100644
index 0000000..afc68df
--- /dev/null
+++ b/test.py
@@ -0,0 +1,10 @@
+import io
+from PIL import Image
+
+img = Image.open('test.jpg', mode='r')
+
+img_byte_arr = io.BytesIO()
+img.save(img_byte_arr, format='PNG')
+img_byte_arr = img_byte_arr.getvalue()
+
+print(img_byte_arr)
\ No newline at end of file