add one function

.idea/misc.xml

+<?xml version="1.0" encoding="UTF-8"?>
+<project version="4">
+  <component name="JavaScriptSettings">
+    <option name="languageLevel" value="ES6" />
+  </component>
+  <component name="ProjectRootManager" version="2" project-jdk-name="Python 3.7" project-jdk-type="Python SDK" />
+</project>
\ No newline at end of file

.idea/modules.xml

+<?xml version="1.0" encoding="UTF-8"?>
+<project version="4">
+  <component name="ProjectModuleManager">
+    <modules>
+      <module fileurl="file://$PROJECT_DIR$/.idea/onnx.iml" filepath="$PROJECT_DIR$/.idea/onnx.iml" />
+    </modules>
+  </component>
+</project>
\ No newline at end of file

.idea/onnx.iml

+<?xml version="1.0" encoding="UTF-8"?>
+<module type="PYTHON_MODULE" version="4">
+  <component name="NewModuleRootManager">
+    <content url="file://$MODULE_DIR$">
+      <sourceFolder url="file://$MODULE_DIR$/onnx_infer/yolo3" isTestSource="false" />
+    </content>
+    <orderEntry type="inheritedJdk" />
+    <orderEntry type="sourceFolder" forTests="false" />
+  </component>
+  <component name="TestRunnerService">
+    <option name="PROJECT_TEST_RUNNER" value="Unittests" />
+  </component>
+</module>
\ No newline at end of file

.idea/vcs.xml

+<?xml version="1.0" encoding="UTF-8"?>
+<project version="4">
+  <component name="VcsDirectoryMappings">
+    <mapping directory="$PROJECT_DIR$" vcs="Git" />
+  </component>
+</project>
\ No newline at end of file

onnx_infer/main.py

@@ -24,8 +24,6 @@ class Main:
         self.colorOnnx_path = colorOnnx_path
         self.yoloOnnx_path = yoloOnnx_path
         self.image_path = image_path
-        self.anchor_path = 'model_data/yolo_anchors.txt'
-        self.classes_path = 'model_data/deepfashion.txt'
         self.score = 0.05
         self.iou = 0.05
         self.picSize = (416, 416)
@@ -40,7 +38,7 @@ class Main:
         feature = []
         for f in precit:
             feature.append(tf.convert_to_tensor(f))
-        sup = Supress(self.anchor_path, self.classes_path, self.score, self.iou, feature)
+        sup = Supress(self.score, self.iou, feature)
         box, score, classes = sup.detect(image)
         with tf.Session() as sess:
             bbox = box.eval()[0]
@@ -58,6 +56,12 @@ class Main:
 
 # main test
 
+def get_result(colorOnnx_path, yoloOnnx_path, image_path):
+    m = Main(colorOnnx_path, yoloOnnx_path, image_path)
+    color, bbox, category = m.colorAndbboxAndcategory()
+    return color, bbox, category
+
+
 if __name__ == '__main__':
     colorOnnx_path = '/Users/apple/Desktop/color.onnx'
     image_path = '/Users/apple/Desktop/8.jpg'

onnx_infer/model_data/deepfashion.txt

-long sleeve dress
-vest dress
-vest
-long sleeve outwear
-long sleeve top
-trousers
-short sleeve top
-sling dress
-skirt
-short sleeve dress
-shorts

onnx_infer/model_data/yolo_anchors.txt

-10,13,  16,30,  33,23,  30,61,  62,45,  59,119,  116,90,  156,198,  373,326

onnx_infer/supression.py

@@ -4,9 +4,7 @@ from yolo3.model import yolo_eval
 import os
 
 class Supress:
-    def __init__(self, anchorPath, classPath, score, iou, featureMap):
-        self.anchorPath = anchorPath
-        self.classPath = classPath
+    def __init__(self, score, iou, featureMap):
         self.score = score
         self.iou = iou
         self.featureMap = featureMap
@@ -14,18 +12,14 @@ class Supress:
 
 
     def _get_class(self):
-        classes_path = os.path.expanduser(self.classPath)
-        with open(classes_path) as f:
-            classNames = f.readlines()
-        classNames = [x.strip() for x in classNames]
+        classNames = ['long sleeve dress', 'vest dress', 'vest', 'long sleeve outwear', 'long sleeve top', 'trousers',
+                      'short sleeve top', 'sling dress', 'skirt', 'short sleeve dress', 'shorts']
         return classNames
 
     def _get_anchors(self):
-        anchorPath = os.path.expanduser(self.anchorPath)
-        with open(anchorPath) as f:
-            anchors = f.readline()
-        anchors = [float(x) for x in anchors.split(',')]
-        return np.array(anchors).reshape(-1, 2)
+        anchors = [[10., 13.], [16., 30.], [33., 23.], [30., 61.], [62., 45.],
+ [59., 119.], [116., 90.], [156., 198.], [373., 326.]]
+        return np.array(anchors)
 
     def detect(self, image):
         anchors = self._get_anchors()

setup.py

+#!/user/bin/python
+
+from setuptools import setup, find_packages
+
+setup(
+    name = 'onnx_infer',
+    version = '0.1.0',
+    description = 'model inference',
+    author = 'wjl',
+    url = 'http://git.wanmeizhensuo.com/wujilong/onnx_model.git',
+    author_email = 'wujilong@igengmei.com',
+    license = 'MIT',
+    packages = ['onnx_infer', 'onnx_infer.yolo3'],
+
+    )
+