├── Dockerfile ├── Makefile ├── README.md ├── model.py └── test.py /Dockerfile: -------------------------------------------------------------------------------- 1 | FROM ubuntu:20.04 2 | 3 | ENV DEBIAN_FRONTEND=noninteractive 4 | 5 | # build-depends of OpenVINO... bit much, could be trimmed down 6 | # didn't include visualization libs 7 | RUN apt update && \ 8 | apt install -y \ 9 | automake \ 10 | autoconf \ 11 | build-essential \ 12 | ca-certificates \ 13 | cmake \ 14 | curl \ 15 | g++-multilib \ 16 | gcc-multilib \ 17 | git \ 18 | git-lfs \ 19 | gstreamer1.0-plugins-base \ 20 | libavcodec-dev \ 21 | libavformat-dev \ 22 | libboost-regex-dev \ 23 | libglib2.0-dev \ 24 | libgstreamer1.0-0 \ 25 | libopenblas-dev \ 26 | libpng-dev \ 27 | libssl-dev \ 28 | libswscale-dev \ 29 | libtool \ 30 | libusb-1.0-0-dev \ 31 | pkg-config \ 32 | python3 \ 33 | python3-pip \ 34 | unzip \ 35 | unzip \ 36 | wget && \ 37 | rm -fr /var/lib/apt/lists 38 | 39 | ARG NUM_JOBS=6 40 | RUN git clone https://github.com/openvinotoolkit/openvino/ --depth=1 --recursive && \ 41 | mkdir -p openvino/build && \ 42 | cd openvino/build && \ 43 | cmake -DCMAKE_BUILD_TYPE=Release -DENABLE_CPPLINT:=OFF -DENABLE_FASTER_BUILD:=ON .. && \ 44 | make -j ${NUM_JOBS} -l ${NUM_JOBS} install && \ 45 | cd ../.. && \ 46 | rm -fr openvino 47 | 48 | RUN pip install -U pip setuptools wheel \ 49 | defusedxml requests networkx \ 50 | torch torchvision onnx \ 51 | openvino 52 | -------------------------------------------------------------------------------- /Makefile: -------------------------------------------------------------------------------- 1 | .PHONY: clean all source abs no_abs docker build inspect 2 | 3 | all: out/model.blob 4 | 5 | clean: 6 | rm out -fr 7 | 8 | SHELL:=bash 9 | OPENVINO_DIR=/usr/local 10 | TOOLS_DIR=${OPENVINO_DIR}/deployment_tools 11 | 12 | BUILDX:=$(shell docker build --help 2>/dev/null | grep -q -- '--push'; echo $$?) 13 | ifeq (${BUILDX},0) 14 | PUSH_ARG=--load 15 | else 16 | PUSH_ARG= 17 | endif 18 | 19 | farid: model.py 20 | python3 model.py --gradient 5 21 | 22 | scharr: model.py 23 | python3 model.py 24 | 25 | out/model.onnx: farid 26 | 27 | out/model.xml: out/model.onnx 28 | python3 ${TOOLS_DIR}/model_optimizer/mo_onnx.py --input_model "out/model.onnx" --data_type half -o out --input_shape "[1, 3, 300, 300]" 29 | 30 | out/model.blob: out/model.xml 31 | . ${OPENVINO_DIR}/bin/setupvars.sh && \ 32 | ${TOOLS_DIR}/tools/compile_tool/compile_tool -m out/model.xml -o out/model.blob -ip U8 -d MYRIAD -VPU_NUMBER_OF_SHAVES 4 -VPU_NUMBER_OF_CMX_SLICES 4 33 | 34 | define DOCKER_CMD_BODY 35 | #! /usr/bin/env sh 36 | 37 | docker run --rm -it -v $(shell pwd):/model -u $(shell id -u) \ 38 | pytorch_harris/openvino:latest "$$@" 39 | endef 40 | export DOCKER_CMD_BODY 41 | out/docker: Dockerfile 42 | docker build -t pytorch_harris/openvino:latest ${PUSH_ARG} -f Dockerfile . 43 | echo "$$DOCKER_CMD_BODY" > out/docker 44 | chmod +x out/docker 45 | 46 | docker: out/docker 47 | 48 | build: out/docker model.py 49 | out/docker bash -c 'cd /model; make out/model.blob' 50 | 51 | inspect: out/docker model.py 52 | out/docker bash 53 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | ## What's this 2 | 3 | A pytorch network (WIP) to run harris corner detection 4 | 5 | ## How to use 6 | 7 | If your device supports pytorch's abs function, use 8 | ``` 9 | model.py --abs 10 | ``` 11 | else use 12 | ``` 13 | model.py 14 | ``` 15 | 16 | ## Using the makefile 17 | 18 | `make build` runs `make docker` + compiles the network in the docker. 19 | 20 | In order to change to `model.py --abs`, you need to manually modify the makefile (I couldn't make it via cli arg). The target `out/model.onnx` needs to require `abs` instead of `no_abs`. 21 | 22 | Sometimes, `make build` needs to be run twice in order to compile the model. 23 | -------------------------------------------------------------------------------- /model.py: -------------------------------------------------------------------------------- 1 | #! /usr/bin/env python3 2 | 3 | import argparse 4 | import numpy as np 5 | from pathlib import Path 6 | import torch 7 | from torch import nn 8 | import torchvision as tv 9 | from torch.nn import functional as F 10 | from typing import Optional, Tuple 11 | 12 | 13 | def get_gaussian_2d(size, mean, stddev): 14 | grid1d = np.linspace(-1, 1, size) 15 | x, y = np.meshgrid(grid1d, grid1d) 16 | grid2d = np.sqrt(x * x + y * y) 17 | 18 | gauss = np.exp(-((grid2d - mean) ** 2) / (2 * stddev ** 2)) 19 | return gauss 20 | 21 | 22 | def get_edge_weights(size, dtype=np.float64): 23 | # For more info about the weights, please see: 24 | # skimage/filters/edges.py 25 | interp = None 26 | smooth = None 27 | 28 | if size == 3: # use scharr 29 | interp = np.matrix([1, 0, -1], dtype=dtype) 30 | smooth = np.matrix([3, 10, 3], dtype=dtype) / 16 31 | elif size == 5: # use farid 32 | interp = np.matrix( 33 | [ 34 | 0.109603762960254, 35 | 0.276690988455557, 36 | 0.0, 37 | -0.276690988455557, 38 | -0.109603762960254, 39 | ], 40 | dtype=dtype, 41 | ) 42 | smooth = np.matrix( 43 | [ 44 | 0.0376593171958126, 45 | 0.249153396177344, 46 | 0.426374573253687, 47 | 0.249153396177344, 48 | 0.0376593171958126, 49 | ], 50 | dtype=dtype, 51 | ) 52 | return interp, smooth 53 | 54 | 55 | def make_edge_filter(interp, smooth, dtype=torch.float32): 56 | x = torch.tensor(smooth.T * interp, dtype=dtype) 57 | y = torch.tensor(interp.T * smooth, dtype=dtype) 58 | # unsqueeze to make this into a filter with shape: [out, in, w, h] 59 | return torch.stack([x, y]).unsqueeze(1) 60 | 61 | 62 | class CornerDetection(nn.Module): 63 | """ 64 | Converts a batch of RGB images into grayscale + blur + edges + Corner detection 65 | 1, C, H, W 66 | """ 67 | 68 | def __init__( 69 | self, 70 | dtype: torch.dtype = torch.float32, 71 | corner_window: int = 3, 72 | nms_window: int = 3, 73 | blur_window: Optional[int] = None, 74 | gradient_window: Optional[int] = None, 75 | gradient_tensor: Optional[torch.Tensor] = None, 76 | ): 77 | """ 78 | gradient_tensor: tensor in [2, 1, K, K] format, of type dtype 79 | If None, a tensor of shape [2, 1, 3, 3] is used by default 80 | 81 | corner_window: Size of window to consider a neighborhood for corner detection 82 | nms_window: Size of window to perform Non-Maximal Supression 83 | 84 | blur_window: 85 | """ 86 | super(CornerDetection, self).__init__() 87 | 88 | self.dtype = dtype 89 | 90 | assert corner_window % 2, "Corner window size needs to be odd to preserve size" 91 | if blur_window: 92 | assert blur_window % 2, "Blur window size needs to be odd to preserve size" 93 | 94 | self.corner_window_size = corner_window 95 | self.nms_window_size = nms_window 96 | self.blur_window_size = blur_window 97 | self.edge_filter = None 98 | 99 | if gradient_tensor is not None: 100 | assert ( 101 | gradient_window is None 102 | ), "Window size not needed when the tensor is provided" 103 | assert gradient_tensor.shape[0] == 2, "Only 2-axis gradients are supported" 104 | assert ( 105 | gradient_tensor.shape[1] == 1 106 | ), "Gradient needs to be batch independent" 107 | assert gradient_tensor.shape[2] % 2, "Gradient window size needs to be odd" 108 | width, height = gradient_tensor.shape[-2:] 109 | assert width == height, "Gradient window size needs to be square" 110 | 111 | assert ( 112 | gradient_tensor.dtype == dtype 113 | ), "Everything needs to be of the same dtype" 114 | 115 | self.edge_filter = gradient_tensor 116 | else: 117 | if gradient_window is None: 118 | gradient_window = 3 119 | assert gradient_window % 2, "Filter size needs to be odd" 120 | self.edge_filter = make_edge_filter( 121 | *get_edge_weights(gradient_window), dtype=self.dtype 122 | ) 123 | 124 | def forward(self, x): 125 | gray = tv.transforms.Grayscale() 126 | image = gray(x) 127 | 128 | edges = F.conv2d( 129 | image, self.edge_filter, padding=int(self.edge_filter.shape[-1] // 2) 130 | ) 131 | 132 | # Shi-Tomasi 133 | # Ix*Ix and Iy*Iy 134 | edges_sqr = edges.multiply(edges) 135 | # Ix*Iy 136 | # edges_xy = edges.prod(dim=1) # Doesn't work with OpenVINO 137 | edges_xy = edges[:, 0, ::] * edges[:, 1, ::] 138 | # Put them together 139 | edges_prod = torch.cat([edges_sqr.squeeze(0), edges_xy]).unsqueeze(0) 140 | # In total 3 layers 141 | conv_dims = 3 142 | # sum all elements in a window 143 | M = F.conv2d( 144 | edges_prod, 145 | padding=self.corner_window_size // 2, 146 | groups=conv_dims, 147 | weight=torch.ones( 148 | conv_dims, 1, self.corner_window_size, self.corner_window_size 149 | ), 150 | ) 151 | # Now we have: 152 | # 0: \sum_{window} IxIx 153 | # 1: \sum_{window} IyIy 154 | # 2: \sum_{window} IxIy 155 | 156 | # min eigen_value = trace/2 - sqrt((trace/2)**2 - determinant) 157 | # trace = m[0]+m[1] 158 | # determinant = m[0]m[1] - m[2]**2 159 | # => eigen_value = (m[0]+m[1])/2 - sqrt(((m[0]-m[1])/2)**2+m[2]**2) 160 | IxIx = M[:, 0, ::] 161 | IyIy = M[:, 1, ::] 162 | IxIy = M[:, 2, ::] 163 | IxIx_IyIy = IxIx - IyIy 164 | trace = IxIx + IyIy 165 | 166 | eig_val = ( 167 | trace - (IxIx_IyIy.multiply_(IxIx_IyIy) + 4 * IxIy.multiply_(IxIy)).sqrt() 168 | ) 169 | # apply non-max suppression, doesn't work on OpenVINO 170 | local_maxima = ( 171 | F.max_pool2d(eig_val, self.nms_window_size, stride=1, padding=1) == eig_val 172 | ) 173 | result = local_maxima * eig_val 174 | 175 | # needs to be (edges**2)**(0.5) because edges.abs() doesn't work with OpenVINO 176 | abs_edges = edges_sqr.sqrt_() 177 | 178 | # used as a short-cut for abs_edges.sum(dim=1).sqrt() 179 | edge = abs_edges.sum(dim=1) / 2 180 | # edge = abs_edges.sum(dim=1).sqrt() 181 | return torch.stack([edge, eig_val]) 182 | 183 | 184 | def export_onnx(args): 185 | """ 186 | Exports the model to an ONNX file. 187 | """ 188 | # Define the expected input shape (dummy input) 189 | shape = (1, 3, 300, 300) 190 | # Create the Model 191 | dtype = torch.float32 192 | model = CornerDetection(gradient_window=args.gradient_window, dtype=dtype) 193 | X = torch.ones(shape, dtype=dtype) 194 | output_file = args.output 195 | print(f"Writing to {output_file}") 196 | torch.onnx.export( 197 | model, 198 | X, 199 | f"{output_file.as_posix()}", 200 | opset_version=12, 201 | do_constant_folding=True, 202 | ) 203 | 204 | 205 | def export(args): 206 | output_dir = args.output.parent 207 | output_dir.mkdir(parents=True, exist_ok=True) 208 | export_onnx(args) 209 | print("Done.") 210 | 211 | 212 | def parse_args(): 213 | parser = argparse.ArgumentParser() 214 | 215 | parser.add_argument( 216 | "-o", "--output", default="out/model.onnx", help="output filename" 217 | ) 218 | parser.add_argument( 219 | "--gradient-window", 220 | type=int, 221 | default=None, 222 | help="Size of window (odd) used for extracting edges", 223 | ) 224 | 225 | args, _ = parser.parse_known_args() 226 | args.output = Path(args.output) 227 | return args 228 | 229 | 230 | if __name__ == "__main__": 231 | args = parse_args() 232 | export(args) 233 | -------------------------------------------------------------------------------- /test.py: -------------------------------------------------------------------------------- 1 | from pathlib import Path 2 | 3 | import cv2 as cv 4 | import numpy as np 5 | 6 | import depthai as dai 7 | 8 | def find_edges(gray): 9 | blurred = cv.GaussianBlur(gray, (3, 3), 0) 10 | 11 | v = np.median(blurred) 12 | sigma = 0.33 13 | 14 | lower = int(max(0, (1.0 - sigma) * v)) 15 | upper = int(min(255, (1.0 + sigma) * v)) 16 | edged = cv.Canny(image, lower, upper) 17 | return edged 18 | 19 | def find_features(img, gray): 20 | dst = cv.cornerHarris(gray,2,3,0.04) 21 | dst = cv.dilate(dst,None) 22 | img[dst>0.01*dst.max()]=[0,0,255] 23 | 24 | return img 25 | 26 | if __name__ == "__main__": 27 | pipeline = dai.Pipeline() 28 | # Source 29 | camera = pipeline.createColorCamera() 30 | camera.setPreviewSize(300, 300) 31 | camera.setResolution(dai.ColorCameraProperties.SensorResolution.THE_1080_P) 32 | camera.setInterleaved(False) 33 | # Ops 34 | detection = pipeline.createNeuralNetwork() 35 | blob_path = Path(__file__).parent / "out" / "model.blob" 36 | detection.setBlobPath(f"{blob_path.as_posix()}") 37 | # Link Camera -> Model 38 | camera.preview.link(detection.input) 39 | # Link Model Output -> Host 40 | x_out = pipeline.createXLinkOut() 41 | x_out.setStreamName("custom") 42 | image = pipeline.createXLinkOut() 43 | image.setStreamName("rgb") 44 | camera.preview.link(image.input) 45 | detection.out.link(x_out.input) 46 | 47 | device = dai.Device(pipeline) 48 | 49 | frame_buffer = device.getOutputQueue(name="custom", maxSize=4) 50 | image = device.getOutputQueue(name="rgb", maxSize=4) 51 | 52 | while True: 53 | frame = frame_buffer.get() 54 | img = image.get().getCvFrame() 55 | gray = cv.cvtColor(img, cv.COLOR_BGR2GRAY) 56 | # Model output 57 | layer = np.array(frame.getFirstLayerFp16()) 58 | # deal with inf 59 | layer[layer == np.Inf] = 255 60 | # Reshape 61 | layer = np.array(layer, dtype=np.uint8) 62 | shape = (2, 300, 300) 63 | frame_data = layer.reshape(shape) 64 | edges = frame_data[0, :,:] * 5 # for visual appeal 65 | corners = frame_data[1, :,:] 66 | 67 | cv.imshow("Image", img) 68 | cv.imshow("Device Corners", corners) 69 | cv.imshow("Device Edges", edges) 70 | 71 | # edged = find_edges(gray) 72 | # cv.imshow("Host Edges", edged) 73 | corner = find_features(img, gray) 74 | cv.imshow("Corner", corner) 75 | 76 | if cv.waitKey(1) == ord("q"): 77 | break 78 | --------------------------------------------------------------------------------