-
Notifications
You must be signed in to change notification settings - Fork 9
/
main.py
190 lines (160 loc) · 6.9 KB
/
main.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
import shutil
import sys
from multiprocessing import Manager
from pathlib import Path
from sanic import Sanic, response
from sanic.config import Config
from sanic.log import logger
from sanic.exceptions import ServerError
import sentry_sdk
from yolo.export_yolov5 import YoloV5Exporter
from yolo.export_yolov6 import YoloV6R4Exporter
from yolo.export_yolov8 import YoloV8Exporter
from yolo.export_yolov10 import YoloV10Exporter
import os
import aiofiles
Config.KEEP_ALIVE = False
Config.RESPONSE_TIMEOUT = 1000
app = Sanic(__name__)
app.config.static_path = (Path(__file__).parent / './client/build/static').resolve().absolute()
if not app.config.static_path.exists():
raise RuntimeError("Client was not built. Please run `npm install && npm run build` to build the client")
app.static('/static', app.config.static_path)
manager = Manager()
conversions = manager.dict()
app.config.workdir = Path(__file__).parent / "tmp"
app.config.workdir.mkdir(exist_ok=True)
app.config.REQUEST_MAX_SIZE = 300_000_000
DEFAULT_NSHAVES = 6
DEFAULT_USE_LEGACY_FRONTEND = 'true'
DEFAULT_USE_RVC2 = 'true'
@app.get("/")
async def index(request):
return await response.file(app.config.static_path / "../index.html")
@app.get("/progress/<key>")
async def index(request, key):
return response.json({"progress": conversions.get(key, "none")})
@app.post('/upload')
async def upload_file(request):
conv_id = str(request.form["id"][0])
logger.info(f"CONVERSION_ID: {conv_id}")
conversions[conv_id] = "new"
nShaves = request.form["nShaves"][0] if "nShaves" in request.form else DEFAULT_NSHAVES
logger.info(f"nShaves: {nShaves}")
useLegacyFrontend = request.form["useLegacyFrontend"][0] if "useLegacyFrontend" in request.form else DEFAULT_USE_LEGACY_FRONTEND
logger.info(f"useLegacyFrontend: {useLegacyFrontend}")
useRVC2 = request.form["useRVC2"][0] if "useRVC2" in request.form else DEFAULT_USE_RVC2
logger.info(f"useRVC2: {useRVC2}")
imgsz = request.form["inputshape"][0]
if " " in imgsz:
imgsz = imgsz.split(" ")
input_shape = [int(imgsz[0]), int(imgsz[1])]
else:
input_shape = int(imgsz)
filename = request.files["file"][0].name
conv_path = app.config.workdir / conv_id
conv_path.mkdir(exist_ok=True)
async with aiofiles.open(conv_path / filename, 'wb') as f:
await f.write(request.files["file"][0].body)
version = request.form["version"][0]
# load exporter and do conversion process
conversions[conv_id] = "read"
try:
sys.path.remove("/app/yolo/yolov5")
except:
pass
if version == "v5":
try:
exporter = YoloV5Exporter(conv_path, filename, input_shape, conv_id, nShaves, useLegacyFrontend, useRVC2)
except ValueError as ve:
sentry_sdk.capture_exception(ve)
raise ServerError(message=str(ve), status_code=518)
except Exception as e:
sentry_sdk.capture_exception(e)
raise ServerError(message="Error while loading model", status_code=520)
elif version == "v6r4":
try:
exporter = YoloV6R4Exporter(conv_path, filename, input_shape, conv_id, nShaves, useLegacyFrontend, useRVC2)
except ValueError as ve:
sentry_sdk.capture_exception(ve)
raise ServerError(message=str(ve), status_code=518)
except Exception as e:
sentry_sdk.capture_exception(e)
raise ServerError(message="Error while loading model (This may be caused by trying to convert older releases 1.0, 2.0 or 3.0, in which case, try to convert using the 'YoloV6 (R1)' or 'YoloV6 (R2, R3)' option).", status_code=516)
elif version == "v8" or version == "v9":
try:
exporter = YoloV8Exporter(conv_path, filename, input_shape, conv_id, nShaves, useLegacyFrontend, useRVC2)
except ValueError as ve:
sentry_sdk.capture_exception(ve)
raise ServerError(message=str(ve), status_code=518)
except Exception as e:
sentry_sdk.capture_exception(e)
if version == "v8":
raise ServerError(message="Error while loading model", status_code=520)
else:
raise ServerError(message="Error while loading model", status_code=515)
elif version == "v10":
try:
exporter = YoloV10Exporter(conv_path, filename, input_shape, conv_id, nShaves, useLegacyFrontend, useRVC2)
except ValueError as ve:
sentry_sdk.capture_exception(ve)
raise ServerError(message=str(ve), status_code=518)
except Exception as e:
sentry_sdk.capture_exception(e)
raise ServerError(message="Error while loading model", status_code=520)
else:
raise ValueError(f"Yolo version {version} is not supported.")
conversions[conv_id] = "initialized"
try:
exporter.export_onnx()
except Exception as e:
sentry_sdk.capture_exception(e)
raise ServerError(message="Error while converting to onnx", status_code=521)
conversions[conv_id] = "onnx"
try:
exporter.export_openvino(version)
except Exception as e:
sentry_sdk.capture_exception(e)
raise ServerError(message="Error while converting to openvino", status_code=522)
conversions[conv_id] = "openvino"
try:
exporter.export_blob()
except Exception as e:
sentry_sdk.capture_exception(e)
raise ServerError(message="Error when exporting to blob, likely due to certain operations being unsupported on RVC3. If interested in further information, please open a GitHub issue.", status_code=526)
# raise ServerError(message="Error while converting to blob", status_code=523)
conversions[conv_id] = "blob"
try:
exporter.export_json()
except Exception as e:
sentry_sdk.capture_exception(e)
raise ServerError(message="Error while making json", status_code=524)
conversions[conv_id] = "json"
try:
zip_file = exporter.make_zip()
except Exception as e:
sentry_sdk.capture_exception(e)
raise ServerError(message="Error while making zip", status_code=525)
conversions[conv_id] = "zip"
return await response.file(
location=zip_file.resolve(),
mime_type="application/zip"
)
@app.on_response
async def cleanup(request, response):
if request.path == "/upload":
conv_id = str(request.form["id"][0])
shutil.rmtree(app.config.workdir / conv_id, ignore_errors=True)
if __name__ == '__main__':
runtime = os.getenv("RUNTIME", "debug")
SENTRY_TOKEN = os.getenv("SENTRY_TOKEN")
logger.info(f"SENTRY_TOKEN: {SENTRY_TOKEN}")
logger.info(f"RVC3_BLOBCONVERTER: {os.getenv('RVC3_BLOBCONVERTER')}")
if SENTRY_TOKEN is not None:
sentry_sdk.init(
dsn=SENTRY_TOKEN
)
if runtime == "prod":
app.run(host="0.0.0.0", access_log=False, workers=8)
else:
app.run(host="0.0.0.0", debug=True, workers=4)