modify input

onnx_infer/main.py

@@ -8,22 +8,22 @@ the result of this is the category and bounding box which is the max score.
 """
 #*******************
 
-from yolodata import path2Img, letterbox_image, path2arr
+from yolodata import path2Img, letterbox_image, path2arr, arr2Img
 from onnx2kera import onnxinfere
 from supression import Supress
 import tensorflow as tf
 from color import COLORS, featureTransform, resize, get_color
 from croppic import cropImage
-
+import cv2
 #global
 categorys = ['long sleeve dress', 'vest dress', 'vest', 'long sleeve outwear', 'long sleeve top',
              'trousers', 'short sleeve top', 'sling dress', 'skirt', 'short sleeve dress', 'shorts']
 
 class Main:
-    def __init__(self, colorOnnx_path, yoloOnnx_path, image_path):
+    def __init__(self, colorOnnx_path, yoloOnnx_path, image_arr):
         self.colorOnnx_path = colorOnnx_path
         self.yoloOnnx_path = yoloOnnx_path
-        self.image_path = image_path
+        self.image_arr = image_arr
         self.score = 0.05
         self.iou = 0.05
         self.picSize = (416, 416)
@@ -32,7 +32,7 @@ class Main:
     def bboxAndcategory(self):
         bbox = []
         category = []
-        image = path2Img(self.image_path)
+        image = arr2Img(self.image_arr)
         image_data = letterbox_image(image, self.picSize)
         precit = onnxinfere(self.yoloOnnx_path, image_data)
         feature = []
@@ -47,7 +47,7 @@ class Main:
 
     def colorAndbboxAndcategory(self):
         bbox, category = self.bboxAndcategory()
-        image_arr = path2arr(self.image_path)
+        image_arr = self.image_arr
         image_crop = cropImage(image_arr, bbox)
         resized = resize(image_crop)
         tmp = featureTransform(resized)
@@ -56,8 +56,8 @@ class Main:
 
 # main test
 
-def get_result(colorOnnx_path, yoloOnnx_path, image_path):
-    m = Main(colorOnnx_path, yoloOnnx_path, image_path)
+def get_result(colorOnnx_path, yoloOnnx_path, image_arr):
+    m = Main(colorOnnx_path, yoloOnnx_path, image_arr)
     color, bbox, category = m.colorAndbboxAndcategory()
     return color, bbox, category
 
@@ -66,6 +66,13 @@ if __name__ == '__main__':
     colorOnnx_path = '/Users/apple/Desktop/color.onnx'
     image_path = '/Users/apple/Desktop/8.jpg'
     yoloOnnx_path = '/Users/apple/Desktop/yolo3.onnx'
+    image_arr = cv2.imread(image_path)
+    color, bbox, category = get_result(colorOnnx_path, yoloOnnx_path, image_arr)
+    print(color)
+    print(bbox)
+    print(category)
+
+"""
     m = Main(colorOnnx_path, yoloOnnx_path, image_path)
     color, bbox, category = m.colorAndbboxAndcategory()
     print('its color is : {}, category is : {} '.format(color, category))
@@ -73,4 +80,4 @@ if __name__ == '__main__':
 #    print(color)
 #    print(bbox)
 #    print(category)
-
+"""

onnx_infer/yolodata.py

@@ -12,6 +12,11 @@ def path2arr(path):
     image = cv2.imread(path, cv2.IMREAD_COLOR)
     return image
 
+def arr2Img(image_arr):
+    image = Image.fromarray(image_arr)
+    return image
+
+# this image is one Image type
 def letterbox_image(image, size):
     iw, ih = image.size
     h, w = size

setup.py

@@ -10,7 +10,9 @@ setup(
     url = 'http://git.wanmeizhensuo.com/wujilong/onnx_model.git',
     author_email = 'wujilong@igengmei.com',
     license = 'MIT',
-    packages = ['onnx_infer', 'onnx_infer.yolo3'],
+    packages = ['onnx_infer.color', onnx_infer.croppic, onnx_infer.drawpic, 'onnx_infer.main',
+                  'onnx_infer.onnx2kera',  'onnx_infer.supression', 'onnx_infer.yolodata', 
+                  'onnx_infer.yolo3.model', 'onnx_infer.yolo3.utils'],
 
     )