Add TjsRawImageProxy and fix the proxying logic to support the depth estimation feature (#30)

whitphx · web-flow · commit 8f5f3888f725 · 2024-01-30T17:56:08.000+09:00
* Add TjsRawImageProxy and fix the proxying logic to support the depth estimation feature

* Support new RawImage() constructor

* Support RawImage.read() without as_url()

* Fix

* Pin Git-hosted resources' revisions
diff --git a/pyodide-e2e/src/tests/pipeline.test.ts b/pyodide-e2e/src/tests/pipeline.test.ts
@@ -6,11 +6,11 @@ suite("transformers.pipeline", () => {
   let pyodide: PyodideInterface;
 
   beforeEach(async () => {
-    pyodide = await setupPyodideForTest();
+    pyodide = await setupPyodideForTest(["numpy", "Pillow"]);
   });
 
   test("zero-shot-image-classification with a local file wrapped by as_url()", async () => {
-    await fetch("https://huggingface.co/spaces/gradio/image_mod/resolve/main/images/lion.jpg")
+    await fetch("https://huggingface.co/spaces/gradio/image_mod/resolve/e07924a/images/lion.jpg")
       .then((response) => response.blob())
       .then((blob) => blob.arrayBuffer())
       .then((arrayBuffer) => {
@@ -38,4 +38,45 @@ result = {item['label']: round(item['score'], 2) for item in data}
     const topLabel = Object.keys(resultObj).reduce((a, b) => resultObj[a] > resultObj[b] ? a : b);
     expect(topLabel).toEqual("lion");
   });
+
+  test("depth-estimation", async () => {
+    await fetch("https://huggingface.co/datasets/Xenova/transformers.js-docs/resolve/db8bd36/bread_small.png")
+      .then((response) => response.blob())
+      .then((blob) => blob.arrayBuffer())
+      .then((arrayBuffer) => {
+        const fileData = new Uint8Array(arrayBuffer);
+        const filePath = "/tmp/bread_small.png";
+        pyodide.FS.writeFile(filePath, fileData);
+      });
+
+    await pyodide.runPythonAsync(`
+from transformers_js_py import import_transformers_js, as_url
+
+transformers = await import_transformers_js()
+
+pipeline = transformers.pipeline
+RawImage = transformers.RawImage
+
+depth_estimator = await pipeline('depth-estimation', 'Xenova/depth-anything-small-hf');
+
+output = await depth_estimator(as_url("/tmp/bread_small.png"))
+`);
+    const outputMap = await pyodide.globals.get("output").toJs()  // Python's dict to JS's Map
+    const output = Object.fromEntries(outputMap);
+
+    const depth = output.depth.toJs();
+    const predictedDepth = output.predicted_depth.toJs();
+
+    // API reference: https://huggingface.co/Xenova/depth-anything-small-hf
+    expect(depth.width).toBe(640)
+    expect(depth.height).toBe(424)
+    expect(depth.channels).toBe(1)
+    expect(predictedDepth).toBeDefined()
+
+    await pyodide.runPythonAsync(`
+output["depth"].save('/tmp/depth.png')
+`);
+    const depthImage: Uint8Array = pyodide.FS.readFile("/tmp/depth.png", { encoding: "binary" });
+    // TODO: How to assert the depth image? Image snapshot is not available in the browser env.
+  })
 });
diff --git a/pyodide-e2e/src/tests/rawimage.test.ts b/pyodide-e2e/src/tests/rawimage.test.ts
@@ -0,0 +1,116 @@
+import type { PyodideInterface } from "pyodide";
+import { beforeEach, describe, it, expect } from "vitest";
+import { setupPyodideForTest } from "./utils";
+
+describe("RawImage", () => {
+  let pyodide: PyodideInterface;
+
+  beforeEach(async () => {
+    pyodide = await setupPyodideForTest(["numpy", "Pillow"]);
+  });
+
+  it("can be initialized via .fromURL()", async () => {
+    await pyodide.runPythonAsync(`
+from transformers_js_py import import_transformers_js
+transformers = await import_transformers_js()
+RawImage = transformers.RawImage
+raw_image = await RawImage.fromURL('https://huggingface.co/datasets/Xenova/transformers.js-docs/resolve/db8bd36/bread_small.png')
+    `)
+    const rawImage = await pyodide.globals.get("raw_image").toJs();
+    expect(rawImage).toBeDefined();
+    expect(rawImage).toHaveProperty("width", 640);
+    expect(rawImage).toHaveProperty("height", 424);
+    expect(rawImage).toHaveProperty("channels", 4);
+  });
+
+  it("can be initialized from a local file via .read()", async () => {
+    await fetch("https://huggingface.co/datasets/Xenova/transformers.js-docs/resolve/db8bd36/bread_small.png")
+      .then((response) => response.blob())
+      .then((blob) => blob.arrayBuffer())
+      .then((arrayBuffer) => {
+        const fileData = new Uint8Array(arrayBuffer);
+        const filePath = "/tmp/bread_small.png";
+        pyodide.FS.writeFile(filePath, fileData);
+      });
+
+    await pyodide.runPythonAsync(`
+from transformers_js_py import import_transformers_js
+transformers = await import_transformers_js()
+
+RawImage = transformers.RawImage
+raw_image = await RawImage.read("/tmp/bread_small.png")
+    `)
+    const rawImage = await pyodide.globals.get("raw_image").toJs();
+    expect(rawImage).toBeDefined();
+    expect(rawImage).toHaveProperty("width", 640);
+    expect(rawImage).toHaveProperty("height", 424);
+    expect(rawImage).toHaveProperty("channels", 4);
+  });
+
+  it("can be initialized via the constructor", async () => {
+    await pyodide.runPythonAsync(`
+from transformers_js_py import import_transformers_js
+transformers = await import_transformers_js()
+
+RawImage = transformers.RawImage
+raw_image = RawImage(bytes([0] * 16*10*3), 16, 10, 3)
+    `)
+
+    const rawImage = await pyodide.globals.get("raw_image").toJs();
+    expect(rawImage).toBeDefined();
+    expect(rawImage).toHaveProperty("width", 16);
+    expect(rawImage).toHaveProperty("height", 10);
+    expect(rawImage).toHaveProperty("channels", 3);
+  });
+
+  it("can be transformed into a numpy array and a PIL image and saved to a local file", async () => {
+    await pyodide.runPythonAsync(`
+from transformers_js_py import import_transformers_js
+transformers = await import_transformers_js()
+RawImage = transformers.RawImage
+raw_image = await RawImage.fromURL('https://huggingface.co/datasets/Xenova/transformers.js-docs/resolve/db8bd36/bread_small.png')
+
+numpy_array = raw_image.to_numpy()
+pil_image = raw_image.to_pil()
+raw_image.save('/tmp/bread_small.png')
+    `)
+    const numpyArrayProxy = await pyodide.globals.get("numpy_array");
+    const numpyArray = numpyArrayProxy.getBuffer("u8");
+    expect(numpyArray.shape).toEqual([424, 640, 4]);
+
+    const pilImage = await pyodide.globals.get("pil_image")
+    expect(pilImage.width).toBe(640);
+    expect(pilImage.height).toBe(424);
+    expect(pilImage.mode).toBe("RGBA");
+  });
+
+  describe("Color transform methods such as .grayscale(), rgb(), and .rgba()", async () => {
+    ([["grayscale", 1], ["rgb", 3], ["rgba", 4]] as const).forEach(([source, sourceChannels]) => {
+      describe(`${source} image`, () => {
+        beforeEach(async () => {
+          await pyodide.runPythonAsync(`
+from transformers_js_py import import_transformers_js
+transformers = await import_transformers_js()
+
+RawImage = transformers.RawImage
+raw_image = RawImage(bytes([0] * 16*10*3), 16, 10, ${sourceChannels})
+`)
+        });
+
+        ([["grayscale", 1], ["rgb", 3], ["rgba", 4]] as const).forEach(([target, targetChannels]) => {
+          it(`can be transformed into a ${target} image`, async () => {
+            await pyodide.runPythonAsync(`
+converted_image = raw_image.${target}()
+`)
+            const convertedImage = await pyodide.globals.get("converted_image").toJs();
+            expect(convertedImage).toBeDefined();
+            expect(convertedImage).toHaveProperty("width", 16);
+            expect(convertedImage).toHaveProperty("height", 10);
+            expect(convertedImage).toHaveProperty("channels", targetChannels);
+          });
+        });
+
+      });
+    });
+  });
+})
diff --git a/pyodide-e2e/src/tests/utils.ts b/pyodide-e2e/src/tests/utils.ts
@@ -5,7 +5,7 @@ import wheelUrl from "transformers-js-py.whl";  // This is the alias from vite.c
 
 export const IS_NODE = typeof window === 'undefined';
 
-export async function setupPyodideForTest(): Promise<PyodideInterface> {
+export async function setupPyodideForTest(requirements: string[] = []): Promise<PyodideInterface> {
   const pyodide = await loadPyodide({
     indexURL: IS_NODE
       ? "node_modules/pyodide"  // pnpm puts pyodide at this path
@@ -25,6 +25,8 @@ export async function setupPyodideForTest(): Promise<PyodideInterface> {
     await micropip.install(wheelUrl);
   }
 
+  await micropip.install(requirements);
+
   await pyodide.runPythonAsync(`
 from transformers_js_py import import_transformers_js
 transformers = await import_transformers_js()
diff --git a/transformers_js_py/__init__.py b/transformers_js_py/__init__.py
@@ -1,4 +1,5 @@
-from typing import Any
+import re
+from typing import Any, Union
 
 import js
 import pyodide.code
@@ -7,6 +8,19 @@
 
 from .url import as_url
 
+try:
+    import numpy as np
+except ImportError:
+    np = None
+
+try:
+    import PIL.Image as PILImage
+except ImportError:
+    PILImage = None
+
+
+rx_class_def_code = re.compile(r"^\s*class\s+([a-zA-Z0-9_]+)\s*{", re.MULTILINE)
+
 
 class TjsModuleProxy:
     def __init__(self, js_obj: pyodide.ffi.JsProxy):
@@ -17,7 +31,7 @@ def __init__(self, js_obj: pyodide.ffi.JsProxy):
     def __getattr__(self, name: str) -> Any:
         res = getattr(self.js_obj, name)
         if isinstance(res, pyodide.ffi.JsProxy):
-            return TjsProxy(res)
+            return proxy_tjs_object(res)
         return res
 
     def __repr__(self) -> str:
@@ -27,18 +41,24 @@ def __repr__(self) -> str:
 class TjsProxy:
     def __init__(self, js_obj: pyodide.ffi.JsProxy):
         self._js_obj = js_obj
+        self._is_class = self._js_obj.typeof == "function" and rx_class_def_code.match(
+            self._js_obj.toString()
+        )  # Ref: https://stackoverflow.com/a/30760236/13103190
 
     def __call__(self, *args: Any, **kwds: Any) -> Any:
-        if hasattr(self._js_obj, "_call"):
-            args = pyodide.ffi.to_js(args)
-            kwds = pyodide.ffi.to_js(kwds)
+        args = pyodide.ffi.to_js(args)
+        kwds = pyodide.ffi.to_js(kwds)
 
+        if hasattr(self._js_obj, "_call"):
             # Transformers.js uses a custom _call() method
             # to make the JS classes callable.
             # https://github.com/xenova/transformers.js/blob/2.4.1/src/utils/core.js#L45-L77
             res = self._js_obj._call(*args, **kwds)
         else:
-            res = self._js_obj(*args, **kwds)
+            if self._is_class:
+                res = self._js_obj.new(*args, **kwds)
+            else:
+                res = self._js_obj(*args, **kwds)
 
         return wrap_or_unwrap_proxy_object(res)
 
@@ -54,17 +74,68 @@ def __setitem__(self, key: Any, value: Any) -> None:
         self._js_obj[key] = value
 
     def __setattr__(self, __name: str, __value: Any) -> None:
-        if __name == "_js_obj":
-            super().__setattr__("_js_obj", __value)
+        if __name == "_js_obj" or __name == "_is_class":
+            super().__setattr__(__name, __value)
         else:
             setattr(self._js_obj, __name, __value)
 
 
+class TjsRawImageClassProxy(TjsProxy):
+    def read(self, input: Union["TjsRawImageProxy", str]):
+        return wrap_or_unwrap_proxy_object(self._js_obj.read(as_url(input)))
+
+
+class TjsRawImageProxy(TjsProxy):
+    def to_numpy(self):
+        if np is None:
+            raise RuntimeError("numpy is not available")
+
+        data = self._js_obj.data  # Uint8ClampedArray|Uint8Array
+        width = self._js_obj.width
+        height = self._js_obj.height
+        channels = self._js_obj.channels
+        return np.asarray(data.to_py()).reshape((height, width, channels))
+
+    def to_pil(self):
+        if PILImage is None:
+            raise RuntimeError("PIL is not available")
+
+        numpy_img = self.to_numpy()
+        if numpy_img.shape[2] == 1:
+            # Gray scale image
+            numpy_img = numpy_img[:, :, 0]
+        return PILImage.fromarray(numpy_img)
+
+    def save(self, path: str):
+        self.to_pil().save(path)
+
+
+def proxy_tjs_object(js_obj: pyodide.ffi.JsProxy):
+    """A factory function that wraps a JsProxy object wrapping a Transformers.js object
+    into a Python object of type TjsProxy or is subclass in the case of a special object
+    such as RawImage.
+    """
+    if js_obj == js._transformers.RawImage:
+        return TjsRawImageClassProxy(js_obj)
+    if js_obj.constructor == js._transformers.RawImage:
+        return TjsRawImageProxy(js_obj)
+    return TjsProxy(js_obj)
+
+
+def to_py_default_converter(value: pyodide.ffi.JsProxy, _ignored1, _ignored2):
+    # Pyodide tries to convert the JS object to a Python object
+    # as best as possible, but it doesn't always work.
+    # In such a case, this custom converter is called
+    # and it wraps the JS object into a TjsProxy object.
+    return proxy_tjs_object(value)
+
+
 def wrap_or_unwrap_proxy_object(obj):
     if isinstance(obj, pyodide.ffi.JsProxy):
         if obj.typeof == "object":
-            return obj.to_py()
-        return TjsProxy(obj)
+            return obj.to_py(default_converter=to_py_default_converter)
+
+        return proxy_tjs_object(obj)
     elif isinstance(obj, pyodide.webloop.PyodideFuture):
         return obj.then(wrap_or_unwrap_proxy_object)
     return obj
