Work on style transfer demos

CMakeLists.txt

@@ -131,8 +131,22 @@ target_include_directories(
 )
 
 
+add_library(bridge_objs STATIC $<TARGET_OBJECTS:bridge>)
+set_target_properties(bridge_objs
+    PROPERTIES ARCHIVE_OUTPUT_DIRECTORY "${CMAKE_BINARY_DIR}"
+)
+
+
 set(BRIDGE_OBJECT_FILES $<TARGET_OBJECTS:bridge>)
 
+# add_custom_command(
+#     TARGET bridge
+#     POST_BUILD
+#     COMMAND ${CMAKE_COMMAND} -E copy_directory
+#             "${CMAKE_CURRENT_SOURCE_DIR}/style-transfer/models"
+#             "$<TARGET_FILE_DIR:StyleTransfer>/style-transfer/models"
+#     COMMENT "NOT! Copying ${PROJECT_ROOT_DIR}/examples/vgg/images to $<TARGET_FILE_DIR:vgg>/images"
+# )
 
 
 
@@ -237,7 +251,8 @@ set(CHAI_LINKER_ARGS
     ${BRIDGE_OBJECT_FILES}
     -L ${LIBTORCH_DIR}/lib
     ${LIBTORCH_LIBS_LINKER_ARGS}
-    --ldflags "-Wl,-rpath,${LIBTORCH_DIR}/lib"
+    --ccflags "-I${BRIDGE_DIR}/include -L${PROJECT_ROOT_DIR}/build"
+    --ldflags "-L${PROJECT_ROOT_DIR}/build -Wl,-rpath,${LIBTORCH_DIR}/lib"
 )
 
 

bridge/include/bridge.h

@@ -15,6 +15,7 @@ typedef double float64_t;
 typedef char bool_t;
 typedef unsigned char uint8_t;
 typedef unsigned int uint32_t;
+typedef unsigned long long uint64_t;
 
 typedef struct bridge_tensor_t {
     float* data;
@@ -24,6 +25,17 @@ typedef struct bridge_tensor_t {
 } bridge_tensor_t;
 
 
+typedef struct bridge_pt_model_t {
+    void* pt_module;
+} bridge_pt_model_t;
+
+typedef struct test_struct_t {
+    int* field;
+} test_struct_t;
+
+
+void hello_world(void);
+
 typedef struct nil_scalar_tensor_t {
     float scalar;
     bridge_tensor_t tensor;
@@ -36,6 +48,14 @@ float* unsafe(const float* arr);
 bridge_tensor_t load_tensor_from_file(const uint8_t* file_path);
 bridge_tensor_t load_tensor_dict_from_file(const uint8_t* file_path,const uint8_t* tensor_key);
 bridge_tensor_t load_run_model(const uint8_t* model_path, bridge_tensor_t input);
+
+bridge_pt_model_t load_model(const uint8_t* model_path);
+
+bridge_tensor_t model_forward(bridge_pt_model_t model, bridge_tensor_t input);
+
+
+bridge_tensor_t model_forward_style_transfer(bridge_pt_model_t model, bridge_tensor_t input);
+
 bridge_tensor_t resize(bridge_tensor_t input,int height,int width);
 bridge_tensor_t imagenet_normalize(bridge_tensor_t input);
 
@@ -107,7 +127,6 @@ proto_bridge_simple(tanhshrink);
 void split_loop(int64_t idx, int64_t n);
 void split_loop_filler(int64_t n,int64_t* ret);
 
-void show_webcam(void);
 
 
 // bridge_tensor_t conv2d(

bridge/lib/bridge.cpp

@@ -16,7 +16,6 @@
 #include <chrono>
 #include <thread>
 
-#include <opencv2/opencv.hpp>
 
 
 #define def_bridge_simple(Name) \
@@ -66,6 +65,19 @@ torch::Tensor bridge_to_torch(bridge_tensor_t &bt) {
     return torch::from_blob(bt.data, shape, torch::kFloat);
 }
 
+torch::Tensor bridge_to_torch(bridge_tensor_t &bt,torch::Device device, bool copy,torch::ScalarType dtype = torch::kFloat32) {
+    std::vector<int64_t> sizes_vec(bt.sizes, bt.sizes + bt.dim);
+    auto shape = torch::IntArrayRef(sizes_vec);
+    auto t = torch::from_blob(bt.data, shape, torch::kFloat);
+    if (device != torch::kCPU)
+        copy = true;
+    if (copy)
+        return t.to(device, dtype, /*non_blocking=*/false, /*copy=*/true);
+    else
+        return t.to(device, dtype, /*non_blocking=*/false, /*copy=*/false);
+
+}
+
 extern "C" float32_t* unsafe(const float32_t* arr) {
     return const_cast<float32_t*>(arr);
 }
@@ -131,6 +143,102 @@ extern "C" bridge_tensor_t load_run_model(const uint8_t* model_path, bridge_tens
     return torch_to_bridge(output);
 }
 
+
+extern "C" bridge_pt_model_t load_model(const uint8_t* model_path) {
+
+    std::cout << "Begin loading model from path: " << model_path << std::endl;
+    std::cout.flush();
+    std::string path(reinterpret_cast<const char*>(model_path));
+    std::cout << "Loading model from path: " << path << std::endl;
+    std::cout.flush();
+
+    try {
+
+        auto* module = new torch::jit::Module(torch::jit::load(path));
+        module->to(torch::kMPS,torch::kFloat16,false);
+        module->eval();
+        std::cout << "Model loaded successfully!" << std::endl;
+        std::cout.flush();
+        return { static_cast<void*>(module) };
+
+        // torch::jit::Module tmp = torch::jit::load(path);
+        // std::cout << "Model loaded successfully!" << std::endl;
+        // std::cout.flush();
+        // auto* module = new torch::jit::Module(std::move(tmp));
+        // std::cout << "Model moved successfully!" << std::endl;
+        // std::cout.flush();
+        // return { static_cast<void*>(module) };
+    } catch (const c10::Error& e) {
+        std::cerr << "error loading the model\n" << e.msg();
+        std::cout << "error loading the model\n" << e.msg();
+        std::cout.flush();
+        std::cerr.flush();
+    }
+    std::cout << "Model loading failed!" << std::endl;
+    std::cout.flush();
+
+    return { nullptr };
+
+
+
+    // bridge_pt_model_t model_wrapper;
+    // torch::jit::Module* pt_module = new torch::jit::Module(); // = (torch::jit::Module*) model_wrapper.pt_module;
+    // try {
+    //     *pt_module = torch::jit::load(mp);
+    //     std::cout << "Model loaded successfully!" << std::endl;
+    //     std::cout.flush();
+    //     model_wrapper.pt_module = pt_module;
+    // } catch (const c10::Error& e) {
+    //     std::cerr << "error loading the model\n" << e.msg();
+    //     std::cout << "error loading the model\n" << e.msg();
+    //     std::cout.flush();
+    //     std::cerr.flush();
+    // }
+
+    // std::cout << pt_module->dump_to_str(false,false,false) << std::endl;
+    // std::cout.flush();
+
+    // return model_wrapper;
+}
+
+
+
+bridge_tensor_t model_forward(bridge_pt_model_t model, bridge_tensor_t input, bool is_vgg_based_model) {
+
+    auto tn_mps = bridge_to_torch(input,torch::kMPS,true,torch::kFloat16);
+    // auto tn_mps = tn.to(torch::kMPS,false,true);
+    auto tn = tn_mps.permute({2, 0, 1}).unsqueeze(0).contiguous();
+
+    std::vector<torch::jit::IValue> ins;
+    ins.push_back(tn);
+
+    auto* module = static_cast<torch::jit::Module*>(model.pt_module);
+    auto o = module->forward(ins).toTensor();
+    auto tn_out = o.squeeze(0).contiguous().permute({1, 2, 0}).contiguous();
+
+    if (is_vgg_based_model) {
+        tn_out = tn_out / 255.0;
+    }
+
+    auto tn_out_cpu = tn_out.to(torch::kCPU,torch::kFloat32,false,true);
+    return torch_to_bridge(tn_out_cpu);
+
+}
+
+extern "C" bridge_tensor_t model_forward(bridge_pt_model_t model, bridge_tensor_t input) {
+    return model_forward(model, input, false);
+}
+
+extern "C" bridge_tensor_t model_forward_style_transfer(bridge_pt_model_t model, bridge_tensor_t input) {
+    return model_forward(model, input, true);
+}
+
+
+extern "C" void hello_world(void) {
+    std::cout << "Hello from C++!" << std::endl;
+    std::cout.flush();
+}
+
 extern "C" bridge_tensor_t increment3(bridge_tensor_t arr) {
     auto t = bridge_to_torch(arr);
     // Increment the tensor
@@ -404,37 +512,37 @@ extern "C" void split_loop_filler(int64_t n,int64_t* ret) {
 
 
 
-cv::VideoCapture open_camera(int cam_index) {
-    cv::VideoCapture cap(cam_index, cv::CAP_AVFOUNDATION);
-    if (!cap.isOpened()) {
-        std::cerr << "Could not open camera index " << cam_index << std::endl;
-        return cv::VideoCapture();
-    }
-    cap.set(cv::CAP_PROP_BUFFERSIZE, 1); // minimal internal buffering
-    cap.set(cv::CAP_PROP_FPS, 60);       // request higher FPS if possible
-    return cap;
-}
+// cv::VideoCapture open_camera(int cam_index) {
+//     cv::VideoCapture cap(cam_index, cv::CAP_AVFOUNDATION);
+//     if (!cap.isOpened()) {
+//         std::cerr << "Could not open camera index " << cam_index << std::endl;
+//         return cv::VideoCapture();
+//     }
+//     cap.set(cv::CAP_PROP_BUFFERSIZE, 1); // minimal internal buffering
+//     cap.set(cv::CAP_PROP_FPS, 60);       // request higher FPS if possible
+//     return cap;
+// }
 
 
-extern "C" void show_webcam(void) {
-    cv::VideoCapture cap;
-    cap = open_camera(0);
+// extern "C" void show_webcam(void) {
+//     cv::VideoCapture cap;
+//     cap = open_camera(0);
 
-    cv::Mat frame_bgr;
+//     cv::Mat frame_bgr;
 
-    while (true) {
-        if (!cap.read(frame_bgr) || frame_bgr.empty()) {
-            std::cerr << "[WARN] Empty frame, exiting" << std::endl;
-            break;
-        }
+//     while (true) {
+//         if (!cap.read(frame_bgr) || frame_bgr.empty()) {
+//             std::cerr << "[WARN] Empty frame, exiting" << std::endl;
+//             break;
+//         }
 
-        cv::imshow("webcam", frame_bgr);
+//         cv::imshow("webcam", frame_bgr);
 
-        if (cv::waitKey(1) == 27) { // ESC key
-            break;
-        }
-    }
+//         if (cv::waitKey(1) == 27) { // ESC key
+//             break;
+//         }
+//     }
 
-    cap.release();
-    cv::destroyAllWindows();
-}
+//     cap.release();
+//     cv::destroyAllWindows();
+// }

demos/video/CMakeLists.txt

@@ -2,6 +2,11 @@
 
 # export OpenCV_DIR="$(brew --prefix opencv)/share/opencv4"
 # /opt/homebrew/opt/opencv/share/opencv4
+
+
+
+
+# This is messy
 find_package(OpenCV 4 REQUIRED)
 
 find_library(ACCELERATE Accelerate REQUIRED)

demos/video/chapel-webcam/.gitignore

@@ -0,0 +1,5 @@
+*.a
+*.so
+*.o
+*.dylib
+lib/savec