完善文档和接口

ILG2021 · Feb 16, 2023 · 5121f5c · 5121f5c
1 parent d4eb71f
commit 5121f5c
Show file tree

Hide file tree

Showing 5 changed files with 52 additions and 22 deletions.
diff --git a/README.md b/README.md
diff --git a/client/client.js b/client/client.js
@@ -0,0 +1,35 @@
+import fetch from "node-fetch";
+
+var url = "https://ultralytics.com/images/zidane.jpg"
+var type = 'jpg'
+var imgData = await fetch(url).then(r => r.buffer()).then(buf => `data:image/${type};base64,`+buf.toString('base64'));
+
+var response = await fetch('http://127.0.0.1:1235/persons',
+    {method: 'POST', body: JSON.stringify({
+        "img": [
+          imgData
+        ]
+      }),
+      "headers": {
+        "content-type": "application/json"
+       }
+    }
+);
+var data = await response.json();
+
+console.log("persons:" , data);
+
+response = await fetch('http://127.0.0.1:1235/objects',
+    {method: 'POST', body: JSON.stringify({
+        "img": [
+          imgData
+        ]
+      }),
+      "headers": {
+        "content-type": "application/json"
+       }
+    }
+);
+data = await response.json();
+
+console.log("objects:" , data);
diff --git a/package.json → client/package.json b/package.json → client/package.json
diff --git a/node_client.js b/node_client.js
diff --git a/yolov5.py b/yolov5.py
@@ -20,7 +20,7 @@ def index():
 
 
 @app.route('/persons', methods=["POST"])
-def analyze():
+def persons():
     if request.method != "POST":
         return
     instances = request.get_json()['img']
@@ -36,8 +36,23 @@ def analyze():
             print(len(output_dict))
     return jsonify(results)
 
+@app.route('/objects', methods=["POST"])
+def objects():
+    if request.method != "POST":
+        return
+    instances = request.get_json()['img']
+    results = []
+    if instances and len(instances) > 0:
+        for instance in instances:
+            encoded_data = instance.split(',')[1]
+            im = Image.open(io.BytesIO(base64.b64decode(encoded_data)))
+            ai_results = models['yolov5s'](im, size=320)  # reduce size=320 for faster inference
+            input_dict = json.loads(ai_results.pandas().xyxy[0].to_json(orient="records"))
+            results.append(input_dict)
+    return json.dumps(results)
+
 if __name__ == "__main__":
-    print("人物检测v1.0")
+    print("对象检测v1.0")
     parser = argparse.ArgumentParser(description="Flask API exposing YOLOv5 model")
     parser.add_argument("--port", default=1235, type=int, help="port number")
     parser.add_argument('--model', nargs='+', default=['yolov5s'], help='model(s) to run, i.e. --model yolov5n yolov5s')