Extend as_url() to deal with PIL images (#67)

* Extend `as_url()` to deal with PIL images

* Fix

Author: whitphx

pyodide-e2e/src/tests/url.test.ts

@@ -0,0 +1,46 @@
+import type { PyodideInterface } from "pyodide";
+import { beforeEach, describe, it, expect } from "vitest";
+import { setupPyodideForTest } from "./utils";
+
+describe("as_url()", () => {
+  let pyodide: PyodideInterface;
+
+  beforeEach(async () => {
+    pyodide = await setupPyodideForTest(["numpy", "Pillow"]);
+  });
+
+  it("converts a file system URL into a URL", async () => {
+    await pyodide.runPythonAsync(`
+from transformers_js_py import as_url
+
+with open("test.bin", "wb") as f:
+    f.write(b"test")
+
+url = as_url("test.bin")
+`);
+    const url = await pyodide.globals.get("url");
+    expect(url).toMatch(/^blob:/);
+  });
+
+  it("converts a bytes variable into a URL", async () => {
+    await pyodide.runPythonAsync(`
+from transformers_js_py import as_url
+
+url = as_url(b"test")
+`);
+    const url = await pyodide.globals.get("url");
+    expect(url).toMatch(/^blob:/);
+  });
+
+  it("converts a PIL image into a URL", async () => {
+    await pyodide.runPythonAsync(`
+from PIL import Image
+from transformers_js_py import as_url
+
+img = Image.new("RGB", (100, 100), "black")
+url = as_url(img)
+`);
+    const url = await pyodide.globals.get("url");
+    expect(url).toMatch(/^blob:/);
+  });
+});

transformers_js_py/url.py

@@ -10,6 +10,14 @@
     warnings.warn("Blob is not available in this environment")
     js_Blob = None
 
+try:
+    import PIL.Image as PILImage
+
+    PILImageImage = PILImage.Image
+except ImportError:
+    PILImage = None  # type: ignore
+    PILImageImage = None  # type: ignore
+
 
 def is_url(url: str) -> bool:
     # Check if the URL is valid, covering all the possible schemes.
@@ -21,11 +29,11 @@ def is_url(url: str) -> bool:
         return False
 
 
-def as_url(data_or_file_path: Union[str, bytes]) -> str:
+def as_url(data_or_file_path: Union[str, bytes, PILImageImage]) -> str:
     """For example, `pipeline('zero-shot-image-classification')`
     requires a URL of the input image file in the browser environment.
-    This function converts a file path on Pyodide's virtual file system
-    to a URL that can be used as the input of such pipelines.
+    This function converts an input data or a input file path on Pyodide's virtual FS
+    into a URL that can be used as the input of such pipelines.
 
     Internally, Transformers.js reads the input with this code:
     https://github.com/xenova/transformers.js/blob/2.6.2/src/utils/image.js#L112-L113
@@ -40,6 +48,12 @@ def as_url(data_or_file_path: Union[str, bytes]) -> str:
             data = f.read()
     elif isinstance(data_or_file_path, bytes):
         data = data_or_file_path
+    elif PILImage and isinstance(data_or_file_path, PILImage.Image):
+        import io
+
+        with io.BytesIO() as f:
+            data_or_file_path.save(f, format="PNG")
+            data = f.getvalue()
     else:
         raise TypeError(
             "data_or_file_path must be str or bytes, "