adding setup.py, imagepy plg

Author: yunlongdong

.gitignore

@@ -1,3 +1,5 @@
 __pycache__
 *.pyc
 .DS_Store
+*egg*
+build/*

plgs/infcnn_plg.py

@@ -0,0 +1,24 @@
+from os.path import join as pjoin
+import numpy as np
+import os.path as osp
+from imagepy.core.engine import Simple
+from imagepy.core.mark import GeometryMark
+from .net import LeNet
+
+path = osp.abspath(osp.dirname(__file__))
+
+class Plugin(Simple):
+    title = 'mnist'
+    note = ['8-bit', 'auto_snap', 'not_channel', 'preview']
+
+    net = LeNet()
+    net.load_mat(pjoin(path, 'lenet.mat'))
+
+    def run(self, ips, imgs, para = None):
+        marks = {'type':'layers', 'body':{}}
+        for i, img in enumerate(imgs):
+            tmp = img.copy().astype('float32')/255.0
+            out = self.net(tmp[None, None, :, :])
+            pred = np.argmax(out, -1)
+            marks['body'][i] = {'type':'text', 'color':(255,0,0), 'body':(2, -2, str(pred))}
+        ips.mark = GeometryMark(marks)

plgs/lenet.mat

@@ -0,0 +1,66 @@
+import numpy as np
+from scipy.io import loadmat
+# infCNN layer
+from infCNN import Conv2d, Dense
+# infCNN op 
+from infCNN import ReLU, Flatten, MaxPool, Softmax 
+from time import time
+
+# op
+relu = ReLU()
+flatten = Flatten()
+maxpool = MaxPool()
+softmax = Softmax()
+
+class LeNet():
+    def __init__(self):
+        self.conv1 = Conv2d(1, 6, 5, 1)
+        self.conv2 = Conv2d(6, 16, 5, 1)
+        self.fc1   = Dense(784, 120)
+        self.fc2   = Dense(120, 84)
+        self.fc3   = Dense(84, 10)
+
+        self.weights = [self.conv1, self.conv2, self.fc1, self.fc2, self.fc3]
+
+    def forward(self, x):
+        out = relu(self.conv1(x))
+        # print(out.shape)
+        out = maxpool(out)
+        # print(out.shape)
+        out = relu(self.conv2(out))
+        # print(out.shape)
+        out = maxpool(out)
+        # print(out.shape)
+        out = flatten(out)
+        # print(out.shape)
+        out = relu(self.fc1(out))
+        # print(out.shape)
+        out = relu(self.fc2(out))
+        # print(out.shape)
+        out = self.fc3(out)
+        # print(out.shape)
+        return softmax(out)
+
+    def __call__(self, x):
+        return self.forward(x)
+
+    def load_mat(self, fn):
+        data = loadmat(fn)
+        w_len = len(data)-4
+        print('load num of weights:', w_len)
+        # minus 4 to skip the header
+        for i in range(w_len//2+1):
+            w = data[str(2*i)]
+            b = data[str(2*i+1)]
+            print(i, w.shape, b.shape)
+            self.weights[i].load_from_torch(data[str(2*i)], data[str(2*i+1)][0, :])
+        print('load done!')
+
+if __name__ == "__main__":
+    net = LeNet()
+    data =  np.ones((1, 1, 28, 28))
+    net.load_mat('train/lenet.mat')
+    start = time()
+    print(net(data))
+    print('infCNN LeNet time:', time()-start)
+

plgs/net.py

@@ -0,0 +1,6 @@
+from setuptools import setup, find_packages
+
+setup(
+    name='infCNN',
+    packages=find_packages()
+)