python: fill out the rest of the bindings and support all models

Acly · Acly · commit 8d283c6ec907 · 2025-12-04T21:49:21.000+01:00
diff --git a/bindings/python/visioncpp/__init__.py b/bindings/python/visioncpp/__init__.py
@@ -1 +1,2 @@
-from .vision import *  # noqa
+from ._lib import Error  # noqa
+from .vision import *  # noqa
diff --git a/bindings/python/visioncpp/_lib.py b/bindings/python/visioncpp/_lib.py
@@ -2,6 +2,7 @@
 import platform
 from pathlib import Path
 from ctypes import c_byte, c_char_p, c_void_p, c_int32, POINTER
+from PIL import Image
 
 
 class Error(Exception):
@@ -23,7 +24,7 @@ def _image_format_to_string(format: int):
 def _image_mode_from_string(mode: str):
     match mode:
         case "RGBA":
-            return 0, 4 # visp::image_format, bytes per pixel
+            return 0, 4  # visp::image_format, bytes per pixel
         case "RGB":
             return 3, 3
         case "L":
@@ -48,17 +49,13 @@ def from_bytes(width: int, height: int, stride: int, format: int, data: bytes):
 
     @staticmethod
     def from_pil_image(image):
-        from PIL import Image
-
         assert isinstance(image, Image.Image), "Expected a PIL Image"
         data = image.tobytes()
         w, h = image.size
         format, bpp = _image_mode_from_string(image.mode)
         return ImageView.from_bytes(w, h, w * bpp, format, data)
 
     def to_pil_image(self):
-        from PIL import Image
-
         mode = _image_format_to_string(self.format)
         size = self.height * self.stride
         data = memoryview((c_byte * size).from_address(self.data))
@@ -145,8 +142,17 @@ def init():
     lib.visp_model_destroy.argtypes = [Model, c_int32]
     lib.visp_model_destroy.restype = None
 
-    lib.visp_esrgan_compute.argtypes = [Model, ImageView, POINTER(ImageView), POINTER(ImageData)]
-    lib.visp_esrgan_compute.restype = c_int32
+    lib.visp_model_compute.argtypes = [
+        Model,
+        c_int32,
+        POINTER(ImageView),
+        c_int32,
+        POINTER(c_int32),
+        c_int32,
+        POINTER(ImageView),
+        POINTER(ImageData),
+    ]
+    lib.visp_model_compute.restype = c_int32
 
     return lib
 
diff --git a/bindings/python/visioncpp/vision.py b/bindings/python/visioncpp/vision.py
@@ -1,7 +1,7 @@
-from ctypes import CDLL, byref
+from ctypes import CDLL, byref, c_int32
 from enum import Enum
 from pathlib import Path
-from typing import NamedTuple
+from typing import NamedTuple, Sequence
 import PIL.Image
 
 from . import _lib as lib
@@ -84,41 +84,58 @@ class Arch(Enum):
     unknown = 5
 
 
-class ESRGAN:
-    arch = Arch.esrgan
-
+class Model:
     @classmethod
-    def load(cls, path: str | Path, device: Device):
+    def load(cls, path: str | Path, device: Device, arch=Arch.unknown):
         api = get_lib()
         handle = lib.Model()
+        path_str = lib.path_to_char_p(path)
+        if arch is Arch.unknown:
+            arch_v = c_int32()
+            check(api.visp_model_detect_family(path_str, byref(arch_v)))
+            arch = Arch(arch_v.value)
+        else:
+            arch_v = arch.value
+
+        check(api.visp_model_load(path_str, device._handle, arch_v, byref(handle)))
+        return cls(api, handle, arch)
+
+    def compute(self, *images: Image, args: Sequence[int] | None = None):
+        if args is None:
+            args = []
+
+        in_views = [_img_view(i) for i in images]
+        in_views_array = (lib.ImageView * len(in_views))(*in_views)
+        args_array = (lib.c_int32 * len(args))(*args)
+        out_view = lib.ImageView()
+        out_data = lib.ImageData()
         check(
-            api.visp_model_load(
-                lib.path_to_char_p(path), device._handle, cls.arch.value, byref(handle)
+            self._api.visp_model_compute(
+                self._handle,
+                self.arch.value,
+                in_views_array,
+                len(in_views_array),
+                args_array,
+                len(args_array),
+                byref(out_view),
+                byref(out_data),
             )
         )
-        return cls(api, handle)
-
-    def compute(self, image: Image):
-        api = self._api
-        in_view = _img_view(image)
-        out_view = lib.ImageView()
-        out_data = lib.ImageData()
-        check(api.visp_esrgan_compute(self._handle, in_view, byref(out_view), byref(out_data)))
         try:
             result = lib.ImageView.to_pil_image(out_view)
         finally:
-            api.visp_image_destroy(out_data)
+            self._api.visp_image_destroy(out_data)
         return result
 
-    def __init__(self, api: CDLL, handle: lib.Handle):
+    def __init__(self, api: CDLL, handle: lib.Handle, arch: Arch):
+        self.arch = arch
         self._api = api
         self._handle = handle
 
     def __del__(self):
         self._api.visp_model_destroy(self._handle, self.arch.value)
 
 
-
 def _img_view(i: Image) -> lib.ImageView:
     if isinstance(i, PIL.Image.Image):
         return lib.ImageView.from_pil_image(i)
diff --git a/src/visp/c-api.cpp b/src/visp/c-api.cpp
@@ -1,7 +1,6 @@
 #include "util/string.h"
 #include "visp/vision.h"
 
-
 using namespace visp;
 
 thread_local fixed_string<512> _error_string{};
@@ -14,20 +13,110 @@ template <typename F>
 int32_t handle_errors(F&& f) {
     try {
         f();
-    } catch (std::exception const& e) {        
+    } catch (std::exception const& e) {
         set_error(e);
         return 0;
     }
     return 1;
 }
 
-extern "C" {
-
-VISP_API char const* visp_get_last_error() {
-    return _error_string.c_str();
+void expect_images(span<image_view> images, size_t count) {
+    if (images.size() != count) {
+        throw except("Expected {} input images, but got {}.", count, images.size());
+    }
 }
 
-// image
+template <model_family f>
+struct model_funcs {};
+
+template <>
+struct model_funcs<model_family::sam> {
+    using model_t = sam_model;
+
+    static sam_model load(char const* filepath, backend_device const& dev) {
+        return sam_load_model(filepath, dev);
+    }
+    static image_data compute(sam_model& m, span<image_view> inputs, span<int> prompt) {
+        expect_images(inputs, 1);
+        sam_encode(m, inputs[0]);
+        if (prompt.size() == 2) {
+            return sam_compute(m, i32x2{prompt[0], prompt[1]});
+        } else if (prompt.size() == 4) {
+            return sam_compute(m, box_2d{i32x2{prompt[0], prompt[1]}, i32x2{prompt[2], prompt[3]}});
+        } else {
+            throw except("sam: bad number of arguments ({}), must be 2 or 4", prompt.size());
+        }
+    }
+};
+
+template <>
+struct model_funcs<model_family::birefnet> {
+    using model_t = birefnet_model;
+
+    static birefnet_model load(char const* filepath, backend_device const& dev) {
+        return birefnet_load_model(filepath, dev);
+    }
+    static image_data compute(birefnet_model& m, span<image_view> inputs, span<int>) {
+        expect_images(inputs, 1);
+        return birefnet_compute(m, inputs[0]);
+    }
+};
+
+template <>
+struct model_funcs<model_family::depth_anything> {
+    using model_t = depthany_model;
+
+    static depthany_model load(char const* filepath, backend_device const& dev) {
+        return depthany_load_model(filepath, dev);
+    }
+    static image_data compute(depthany_model& m, span<image_view> inputs, span<int>) {
+        expect_images(inputs, 1);
+        image_data result_f32 = depthany_compute(m, inputs[0]);
+        image_data normalized = image_normalize(result_f32);
+        return image_f32_to_u8(normalized, image_format::alpha_u8);
+    }
+};
+
+template <>
+struct model_funcs<model_family::migan> {
+    using model_t = migan_model;
+
+    static migan_model load(char const* filepath, backend_device const& dev) {
+        return migan_load_model(filepath, dev);
+    }
+    static image_data compute(migan_model& m, span<image_view> inputs, span<int>) {
+        expect_images(inputs, 2);
+        if (inputs[1].format != image_format::alpha_u8) {
+            throw except("migan: second input image (mask) must be alpha_u8 format");
+        }
+        return migan_compute(m, inputs[0], inputs[1]);
+    }
+};
+
+template <>
+struct model_funcs<model_family::esrgan> {
+    using model_t = esrgan_model;
+
+    static esrgan_model load(char const* filepath, backend_device const& dev) {
+        return esrgan_load_model(filepath, dev);
+    }
+    static image_data compute(esrgan_model& m, span<image_view> inputs, span<int>) {
+        expect_images(inputs, 1);
+        return esrgan_compute(m, inputs[0]);
+    }
+};
+
+template <typename F>
+void dispatch_model(model_family family, F&& f) {
+    switch (family) {
+        case model_family::sam: f(model_funcs<model_family::sam>{}); break;
+        case model_family::birefnet: f(model_funcs<model_family::birefnet>{}); break;
+        case model_family::depth_anything: f(model_funcs<model_family::depth_anything>{}); break;
+        case model_family::migan: f(model_funcs<model_family::migan>{}); break;
+        case model_family::esrgan: f(model_funcs<model_family::esrgan>{}); break;
+        default: throw visp::exception("Unsupported model family");
+    }
+}
 
 struct visp_image_view {
     int32_t width;
@@ -50,6 +139,17 @@ void return_image(image_data** out_data, visp_image_view* out_image, image_data&
     put_image(out_image, **out_data);
 }
 
+//
+// public C interface
+
+extern "C" {
+
+VISP_API char const* visp_get_last_error() {
+    return _error_string.c_str();
+}
+
+// image
+
 VISP_API void visp_image_destroy(image_data* img) {
     delete img;
 }
@@ -107,55 +207,41 @@ VISP_API int32_t visp_model_load(
             model_file file = model_load(filepath);
             family = model_detect_family(file);
         }
-        switch (family) {
-            case model_family::sam: {
-                sam_model model = sam_load_model(filepath, *dev);
-                *out = reinterpret_cast<any_model*>(new sam_model(std::move(model)));
-                break;
-            }
-            case model_family::birefnet: {
-                birefnet_model model = birefnet_load_model(filepath, *dev);
-                *out = reinterpret_cast<any_model*>(new birefnet_model(std::move(model)));
-                break;
-            }
-            case model_family::depth_anything: {
-                depthany_model model = depthany_load_model(filepath, *dev);
-                *out = reinterpret_cast<any_model*>(new depthany_model(std::move(model)));
-                break;
-            }
-            case model_family::migan: {
-                migan_model model = migan_load_model(filepath, *dev);
-                *out = reinterpret_cast<any_model*>(new migan_model(std::move(model)));
-                break;
-            }
-            case model_family::esrgan: {
-                esrgan_model model = esrgan_load_model(filepath, *dev);
-                *out = reinterpret_cast<any_model*>(new esrgan_model(std::move(model)));
-                break;
-            }
-            default: throw visp::exception("Invalid model family");
-        }
+        dispatch_model(family, [&](auto funcs) {
+            using model_t = typename decltype(funcs)::model_t;
+            *out = reinterpret_cast<any_model*>(new model_t(funcs.load(filepath, *dev)));
+        });
     });
 }
 
 VISP_API void visp_model_destroy(any_model* model, int32_t arch) {
     model_family family = model_family(arch);
-    switch (family) {
-        case model_family::sam: delete reinterpret_cast<sam_model*>(model); break;
-        case model_family::birefnet: delete reinterpret_cast<birefnet_model*>(model); break;
-        case model_family::depth_anything: delete reinterpret_cast<depthany_model*>(model); break;
-        case model_family::migan: delete reinterpret_cast<migan_model*>(model); break;
-        case model_family::esrgan: delete reinterpret_cast<esrgan_model*>(model); break;
-        default: fprintf(stderr, "Invalid model family: %d\n", int(family)); break;
-    }
+    dispatch_model(family, [&](auto funcs) {
+        using model_t = typename decltype(funcs)::model_t;
+        delete reinterpret_cast<model_t*>(model);
+    });
 }
 
-VISP_API int32_t visp_esrgan_compute(
-    esrgan_model* model, image_view in_image, visp_image_view* out_image, image_data** out_data) {
+VISP_API int32_t visp_model_compute(
+    any_model* model,
+    int32_t family,
+    image_view* inputs,
+    int32_t n_inputs,
+    int32_t* args,
+    int32_t n_args,
+    visp_image_view* out_image,
+    image_data** out_data) {
 
     return handle_errors([&]() {
-        image_data result = esrgan_compute(*model, in_image);
-        return_image(out_data, out_image, std::move(result));
+        span<image_view> input_views(inputs, n_inputs);
+        span<int32_t> input_args(args, n_args);
+
+        dispatch_model(model_family(family), [&](auto funcs) {
+            using model_t = typename decltype(funcs)::model_t;
+            model_t& m = *reinterpret_cast<model_t*>(model);
+            image_data result = funcs.compute(m, input_views, input_args);
+            return_image(out_data, out_image, std::move(result));
+        });
     });
 }
 
diff --git a/tests/test_python_bindings.py b/tests/test_python_bindings.py

Original file line number	Diff line number	Diff line change
`@@ -1 +1,2 @@`
`1`		`-from .vision import * # noqa`
	`1`	`+from ._lib import Error # noqa`
	`2`	`+from .vision import * # noqa`