add flask

bs7280 · Mar 10, 2022 · e91d3c3 · e91d3c3
1 parent 86a5e1c
commit e91d3c3
Show file tree

Hide file tree

Showing 7 changed files with 144 additions and 6 deletions.
diff --git a/docs/vgg000.png b/docs/vgg000.png
diff --git a/docs/vgg001.png b/docs/vgg001.png
diff --git a/docs/vgg002.png b/docs/vgg002.png
diff --git a/inference.py b/inference.py
@@ -86,12 +86,12 @@ def detect3d(
 
         # Run detection 2d
         dets = detect2d(
-            weights=opt.weights,
+            weights='yolov5s.pt',
             source=img_path,
-            data=opt.data,
-            imgsz=opt.imgsz,
-            device=opt.device,
-            classes=opt.classes
+            data='data/coco128.yaml',
+            imgsz=[640, 640],
+            device=0,
+            classes=[0, 2, 3, 5]
         )
 
         for det in dets:

diff --git a/main.py b/main.py
@@ -0,0 +1,89 @@
+import os
+from flask import Flask, render_template, request
+import cv2
+import numpy as np
+import base64
+from werkzeug.utils import secure_filename
+import shutil
+from pathlib import Path
+import sys
+import argparse
+
+from inference import detect3d
+
+FILE = Path(__file__).resolve()
+ROOT = FILE.parents[0]  # YOLOv5 root directory
+if str(ROOT) not in sys.path:
+    sys.path.append(str(ROOT))  # add ROOT to PATH
+ROOT = Path(os.path.relpath(ROOT, Path.cwd()))  # relative
+
+app = Flask(__name__)
+
+UPLOAD_FOLDER = os.path.basename('static')
+app.config['UPLOAD_FOLDER'] = UPLOAD_FOLDER
+
+@app.route("/")
+def start_page():
+    print("Start")
+    return render_template('index.html')
+
+@app.route('/upload', methods=['POST'])
+def upload_file():
+    FILENAME = {}
+    # get file 
+    image = request.files['image']
+
+    # save file
+    image.save('static/image_eval.png')
+
+    if 'image' in request.files:
+        detect = True
+
+        # process file
+        detect3d(
+            reg_weights='weights/epoch_10.pkl',
+            model_select='resnet',
+            source='static',
+            calib_file='eval/camera_cal/calib_cam_to_cam.txt',
+            save_result=True,
+            show_result=False,
+            output_path='static/'
+        )
+
+        # encode to base64 image
+        with open('static/000.png', 'rb') as image_file:
+            img_encode = base64.b64encode(image_file.read())
+            to_send = 'data:image/png;base64, ' + str(img_encode, 'utf-8')
+    else:
+        detect = False
+
+    return render_template('index.html', init=True, detect=detect, image_to_show=to_send)
+
+def parse_opt():
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--weights', nargs='+', type=str, default=ROOT / 'yolov5s.pt', help='model path(s)')
+    parser.add_argument('--source', type=str, default=ROOT / 'eval/image_2', help='file/dir/URL/glob, 0 for webcam')
+    parser.add_argument('--data', type=str, default=ROOT / 'data/coco128.yaml', help='(optional) dataset.yaml path')
+    parser.add_argument('--imgsz', '--img', '--img-size', nargs='+', type=int, default=[640], help='inference size h,w')
+    parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
+    parser.add_argument('--classes', default=[0, 2, 3, 5], nargs='+', type=int, help='filter by class: --classes 0, or --classes 0 2 3')
+    parser.add_argument('--reg_weights', type=str, default='weights/epoch_10.pkl', help='Regressor model weights')
+    parser.add_argument('--model_select', type=str, default='resnet', help='Regressor model list: resnet, vgg, eff')
+    parser.add_argument('--calib_file', type=str, default=ROOT / 'eval/camera_cal/calib_cam_to_cam.txt', help='Calibration file or path')
+    parser.add_argument('--show_result', action='store_true', help='Show Results with imshow')
+    parser.add_argument('--save_result', action='store_true', help='Save result')
+    parser.add_argument('--output_path', type=str, default=ROOT / 'output', help='Save output pat')
+
+    opt = parser.parse_args()
+    opt.imgsz *= 2 if len(opt.imgsz) == 1 else 1  # expand
+    return opt
+
+
+if __name__ == '__main__':
+    try:
+        os.makedirs('static')
+    except:
+        print('Directory already exist!')
+    opt = parse_opt()
+    app.run()
+    shutil.rmtree('static')
diff --git a/templates/index.html b/templates/index.html
@@ -0,0 +1,49 @@
+<!DOCTYPE html>
+<html lang="en">
+    <head>
+        <meta charset="UTF-8">
+        <title>YOLO3D</title>
+        <link rel="stylesheet" href="https://maxcdn.bootstrapcdn.com/bootstrap/3.3.7/css/bootstrap.min.css" integrity="sha384-BVYiiSIFeK1dGmJRAkycuHAHRg32OmUcww7on3RYdg4Va+PmSTsz/K68vbdEjh4u" crossorigin="anonymous">
+    </head>
+
+    <body>
+        <div class="container">
+            <div class="starter-template">
+                <h1>3D Object Detection with YOLO</h1>
+                <p>By Didi Ruhyadi - github.com/ruhyadi/yolo3d</p>
+            </div>
+        </div>
+
+        <header class="navbar">
+            <div align="center">
+                <h4>Upload Image</h4>
+                <form action="/upload" method="post" enctype="multipart/form-data" >
+                    <div align="center">
+                        <label class="btn btn-default btn-file">
+                            Browse Image<input type="file" name="image" style="display: none;">
+                        </label>
+                        <input type="submit" value="Submit" class="btn btn-primary">
+                    </div>
+                </form>
+            </div>
+
+        </header>
+
+        <div class="text-center">
+            {% if init %}
+                {% if detect %}
+                    <div class="alert alert-success" style="margin-top:18px;">
+                        Pattern <strong>Detected</strong> on Video
+                    </div>
+                    <div>
+                        <img src="{{ image_to_show }}">
+                    </div>
+                {% else %}
+                    <div class="alert alert-danger" style="margin-top:18px;">
+                        No Pattern Detected
+                    </div>
+                {% endif %}
+            {% endif %}
+        </div>
+    </body>
+</html>
diff --git a/weights/get_weights.py b/weights/get_weights.py
@@ -11,7 +11,7 @@ def get_weights(weights):
     weights_list = {
         'resnet': '1Bw4gUsRBxy8XZDGchPJ_URQjbHItikjw',
         'resnet18': '1k_v1RrDO6da_NDhBtMZL5c0QSogCmiRn',
-        'vgg11': ''
+        'vgg11': '1vZcB-NaPUCovVA-pH-g-3NNJuUA948ni'
     }
 
     url = f"https://drive.google.com/uc?id={weights_list[weights]}"