Merge branch 'main' into im-dev

Author: Iainmon

bridge/include/bridge.h

@@ -40,6 +40,8 @@ bridge_tensor_t resize(bridge_tensor_t input,int height,int width);
 bridge_tensor_t imagenet_normalize(bridge_tensor_t input);
 
 
+bridge_tensor_t add_two_arrays(bridge_tensor_t a, bridge_tensor_t b);
+
 // bridge_tensor_t capture_webcam_bridge(int cam_index);
 
 int baz(void);

bridge/lib/bridge.cpp

@@ -264,6 +264,16 @@ extern "C" bridge_tensor_t imagenet_normalize(bridge_tensor_t input) {
     return torch_to_bridge(output);
 }
 
+
+extern "C" bridge_tensor_t add_two_arrays(bridge_tensor_t a, bridge_tensor_t b) {
+    torch::Tensor t_a = bridge_to_torch(a);
+    torch::Tensor t_b = bridge_to_torch(b);
+
+    torch::Tensor output = t_a + t_b;
+
+    return torch_to_bridge(output);
+}
+
 // extern "C" bridge_tensor_t capture_webcam_bridge(int cam_index) {
 //     torch::Tensor image = capture_webcam(cam_index);
 //     return torch_to_bridge(image);

demos/CMakeLists.txt

@@ -1,4 +1,4 @@
-# add_subdirectory(video)
+add_subdirectory(video)
 
 # add_subdirectory(webcam_filter)
 

demos/video/CMakeLists.txt

@@ -11,7 +11,7 @@ find_library(FOUNDATION Foundation REQUIRED)
 
 add_executable(VidStreamer
     ${CMAKE_CURRENT_SOURCE_DIR}/webcam_infer.cpp
-    # ${CMAKE_CURRENT_SOURCE_DIR}/cvutil.hpp
+    ${CMAKE_CURRENT_SOURCE_DIR}/cvtool.hpp
     ${CMAKE_CURRENT_SOURCE_DIR}/imageops.hpp
 )
 
@@ -37,7 +37,6 @@ target_link_libraries(VidStreamer
         ${ACCELERATE}
         ${METAL}
         ${FOUNDATION}
-        BridgeUtil
 )
 
 set_target_properties(VidStreamer PROPERTIES

demos/video/cvtool.hpp

@@ -0,0 +1,235 @@
+#pragma once
+
+#include <chrono>
+#include <iostream>
+#include <opencv2/opencv.hpp>
+#include <torch/script.h>
+#include <torch/torch.h>
+#include <utility>
+
+
+static torch::Device default_device(torch::kCPU);
+torch::Device get_default_device();
+
+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(const std::string &file_path) {
+    cv::VideoCapture cap(file_path);
+    if (!cap.isOpened()) {
+        std::cerr << "Could not open file " << file_path << 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;
+}
+
+std::pair<int,int> get_camera_resolution(cv::VideoCapture& cap) {
+    int width = static_cast<int>(cap.get(cv::CAP_PROP_FRAME_WIDTH));
+    int height = static_cast<int>(cap.get(cv::CAP_PROP_FRAME_HEIGHT));
+    return {height, width};
+}
+
+std::shared_ptr<cv::Mat> create_frame_buffer(int height, int width) {
+    auto frame_buffer = std::make_shared<cv::Mat>(height, width, CV_8UC3);
+    return frame_buffer;
+}
+
+
+
+// std::shared_ptr<at::Tensor> get_frame_buffer_tensor(int height,int width) {
+//     auto options_cpu  = torch::TensorOptions().dtype(torch::kFloat32).device(torch::kCPU);
+//     torch::Tensor frame_tensor_cpu = torch::empty({1, height, width, 3}, options_cpu);  
+// }
+
+std::shared_ptr<at::Tensor> create_buffer_tensor(
+    torch::IntArrayRef sizes,
+    torch::ScalarType = torch::kFloat32,
+    torch::Device device = get_default_device()) {
+    auto options_device = torch::TensorOptions()
+                        .dtype(torch::kFloat32)
+                        .device(default_device);
+    auto tensor = torch::empty(sizes, options_device);
+    auto frame_tensor_device = std::make_shared<at::Tensor>(tensor);
+    return frame_tensor_device;
+}
+
+
+std::shared_ptr<at::Tensor> create_frame_buffer_tensor(int height,int width,torch::Device device = get_default_device()) {
+    torch::IntArrayRef sizes = {1, height, width, 3};
+    return create_buffer_tensor(sizes, torch::kFloat32);
+}
+
+at::Tensor to_tensor(cv::Mat &img) {
+    auto t = torch::from_blob(img.data, {1, img.rows, img.cols, 3}, torch::kUInt8).clone();
+    t = t.to(default_device);
+    t = t.to(torch::kFloat32).permute({0, 3, 1, 2}) / 255.0;
+    return t;//.to(default_device,true);
+}
+
+cv::Mat to_mat(at::Tensor &tensor) {
+    // Ensure the tensor is on the CPU and not on the GPU
+    // at::Tensor cpu_tensor = tensor.to(torch::kCPU);
+
+    // Clone the tensor to avoid modifying the original data
+    // at::Tensor cloned_tensor = cpu_tensor.clone();
+
+    
+    int height = tensor.size(2);
+    int width = tensor.size(3);
+    auto t = tensor
+                .mul(255)
+                .squeeze()
+                .detach()
+                .permute({1, 2, 0})
+                .contiguous()
+                .to(torch::kUInt8)
+                // .clamp(0, 255)
+                .clone()
+                .to(torch::kCPU);
+    cv::Mat mat = cv::Mat(height, width, CV_8UC3, t.data_ptr());
+    return mat;
+
+
+
+    // tensor = tensor.squeeze().detach();
+    // tensor = tensor.permute({1, 2, 0}).contiguous();
+    // tensor = tensor.mul(255).clamp(0, 255).to(torch::kU8);
+    // tensor = tensor.to(torch::kCPU);
+    // int64_t height = tensor.size(0);
+    // int64_t width = tensor.size(1);
+    // cv::Mat mat =
+    //     cv::Mat(cv::Size(width, height), CV_8UC3, tensor.data_ptr<uchar>());
+    // return mat.clone();
+}
+
+torch::Device get_default_device() {
+    if (torch::mps::is_available()) {
+        // default_device = torch::Device(torch::kMPS);
+        std::cout << "[INFO] Running on MPS" << std::endl;
+    } else {
+        std::cout << "[INFO] MPS not available, falling back to CPU" << std::endl;
+    }
+    return default_device;
+}
+
+// torch::jit::Module load_module_from_file(std::string model_path) {
+//     std::string mp(reinterpret_cast<const char*>(model_path));
+
+//     std::cout << "Loading model from path: " << mp << std::endl;
+//     std::cout.flush();
+
+//     torch::jit::Module module;
+//     try
+//     {
+//         // Deserialize the ScriptModule from a file using torch::jit::load().
+//         module = torch::jit::load(mp);
+//     }
+//     catch (const c10::Error& e)
+//     {
+//         std::cerr << "error loading the model\n" << e.msg();
+//         std::system("pause");
+//     }
+
+//     std::vector<torch::jit::IValue> inputs;
+//     inputs.push_back(t_input);
+
+//     return module;
+// }
+
+at::Tensor imagenet_resize(at::Tensor& image, int height, int width) {
+    // Resize the image to the specified height and width
+    auto resized_image = torch::nn::functional::interpolate(
+        image,
+        torch::nn::functional::InterpolateFuncOptions()
+            .size(std::vector<int64_t>({height, width}))
+            .mode(torch::kBilinear)
+            .align_corners(false)
+    );
+    return resized_image;
+}
+
+at::Tensor imagenet_normalize_tensor(at::Tensor& input) {
+    // Normalize the image using ImageNet mean and std
+    // auto mean = torch::tensor({0.485, 0.456, 0.406}).view({1, 3, 1, 1});
+    // auto std = torch::tensor({0.229, 0.224, 0.225}).view({1, 3, 1, 1});
+    // return (image - mean) / std;
+
+    // std::cout << "Input sizes: " << input.sizes() << std::endl;
+
+    at::Tensor image = input.to(torch::kFloat32).clone();// / 255.0;
+    // std::cout << "Image sizes: " << image.sizes() << std::endl;
+
+    static const std::vector<float> mean_data{0.485, 0.456, 0.406};
+    static const std::vector<float> std_data{0.229, 0.224, 0.225};
+    auto options = image.options();
+    auto mean = torch::tensor(mean_data,options).reshape({3, 1, 1});  // (3,1,1)
+    auto std  = torch::tensor(std_data,options).reshape({3, 1, 1});
+
+    if (image.dim() == 4) {
+        mean = mean.unsqueeze(0); // (1,3,1,1)
+        std = std.unsqueeze(0);
+    }
+
+    // std::cout << "Mean sizes: " << mean.sizes() << std::endl;
+    // std::cout << "Std sizes: " << std.sizes() << std::endl;
+    // std::cout << "Image sizes: " << image.sizes() << std::endl;
+    // std::exit(0);
+
+    auto output = (image - mean) / std;
+    output = output;
+    // std::cout << "Output sizes: " << output.sizes() << std::endl;
+    return output;
+}
+
+
+int show_webcam(int cam_index) {
+    cv::VideoCapture cap = open_camera(cam_index);
+    if (!cap.isOpened()) {
+        std::cerr << "Could not open camera index " << cam_index << std::endl;
+        return -1;
+    }
+
+    cv::Mat frame;
+    while (true) {
+        cap >> frame;
+        if (frame.empty()) {
+            std::cerr << "Failed to capture image from camera" << std::endl;
+            break;
+        }
+
+        cv::imshow("Webcam", frame);
+        if (cv::waitKey(30) >= 0) break; // Exit on any key press
+    }
+    return 0;
+}
+
+
+
+at::Tensor capture_webcam(int cam_index) {
+    cv::VideoCapture cap = open_camera(cam_index);
+    if (!cap.isOpened()) {
+        std::cerr << "Could not open camera index " << cam_index << std::endl;
+        return at::Tensor();
+    }
+
+    cv::Mat frame;
+    cap >> frame;
+
+    if (frame.empty()) {
+        std::cerr << "Failed to capture image from camera" << std::endl;
+        return at::Tensor();
+    }
+
+    auto tensor = to_tensor(frame);
+    return tensor;
+}

demos/video/webcam_infer.cpp

@@ -5,7 +5,7 @@
 #include <chrono>
 #include <utility>
 
-#include <cvtool.hpp>
+#include "cvtool.hpp"
 #include "imageops.hpp"
 
 
@@ -54,8 +54,8 @@ struct Model : torch::nn::Module {
     }
     // auto output = x + r;
     auto input = x;
-    // auto output = imageops::sobel_rgb(input);
-    auto output = imagenet_normalize_tensor(input);
+    auto output = imageops::sobel_rgb(input);
+    // auto output = imagenet_normalize_tensor(input);
     return output;
   }
 

lib/Bridge.chpl

@@ -83,6 +83,10 @@ module Bridge {
     extern "imagenet_normalize" proc imageNetNormalize(
         in input: bridge_tensor_t): bridge_tensor_t;
 
+    extern "add_two_arrays" proc addTwoArrays(
+        in a: bridge_tensor_t, 
+        in b: bridge_tensor_t): bridge_tensor_t;
+
     // extern "capture_webcam_bridge" proc captureWebcam(
     //     in cam_index: int(32)): bridge_tensor_t;
 

lib/NDArray.chpl

@@ -1881,6 +1881,13 @@ proc type ndarray.maxPool2d(
     ) : ndarray(inputRank,eltType);
 }
 
+proc type ndarray.addTwoArrays(a: ndarray(?rank,?eltType),b: ndarray(rank,eltType)): ndarray(rank,eltType) {
+    return Bridge.addTwoArrays(
+        a : Bridge.tensorHandle(eltType),
+        b : Bridge.tensorHandle(eltType)
+    ) : ndarray(rank,eltType);
+}
+
 
 proc type ndarray.maxPool(features: ndarray(3, ?eltType), poolSize: int) do
     return this.maxPool(features,poolSize,poolSize);

test/CMakeLists.txt

@@ -1,3 +1,4 @@
 
 
-add_subdirectory(tiny)
+add_subdirectory(tiny)
+add_subdirectory(ndarray_addTwoArrays)

test/ndarray_addTwoArrays/CMakeLists.txt

@@ -0,0 +1,14 @@
+add_executable(AddTwoArraysTest 
+    ${CMAKE_CURRENT_SOURCE_DIR}/addTwoArraysTest.chpl
+    ${CHAI_LIB_FILES}
+)
+add_dependencies(AddTwoArraysTest bridge)
+add_dependencies(AddTwoArraysTest ChAI)
+target_link_options(AddTwoArraysTest
+    PRIVATE
+    ${CHAI_LINKER_ARGS}
+)
+
+set_target_properties(AddTwoArraysTest PROPERTIES
+    RUNTIME_OUTPUT_DIRECTORY ${CMAKE_BINARY_DIR}
+)