forked from ultralytics/yolov5
-
Notifications
You must be signed in to change notification settings - Fork 0
/
server.py
111 lines (89 loc) · 3.71 KB
/
server.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
from http.server import ThreadingHTTPServer, BaseHTTPRequestHandler
import json
import tempfile
import os
import cgi
import time
import subprocess
from detect import run
# Define the request handler class
class SimpleHTTPRequestHandler(BaseHTTPRequestHandler):
# Handle GET requests
def do_GET(self):
self.send_response(200) # Send 200 OK status code
self.send_header("Content-type", "text/html")
self.end_headers()
# Send the HTML form as the response body
form_html = """
<html>
<body>
<form method="POST" enctype="multipart/form-data">
<input type="file" name="file" />
<input type="submit" value="Upload" />
</form>
</body>
</html>
"""
self.wfile.write(form_html.encode("utf-8"))
# Handle POST requests
def do_POST(self):
content_type = self.headers["Content-Type"]
reqdir = "/app/tmp/" + str(time.time()) + "/"
os.makedirs(reqdir, exist_ok=True)
# Check if the content type is multipart/form-data
if content_type.startswith("multipart/form-data"):
# Parse the form data
form_data = cgi.FieldStorage(
fp=self.rfile, headers=self.headers, environ={"REQUEST_METHOD": "POST"}
)
# Get the uploaded file field
file_field = form_data["file"]
# Create a temporary file to save the uploaded file
with tempfile.NamedTemporaryFile(dir="/app/tmp", delete=False) as tmp_file:
# Save the file data to the temporary file
tmp_file.write(file_field.file.read())
# Get the temporary file path
tmp_file_path = tmp_file.name
# Extract the filename from the uploaded file field
filename = os.path.basename(file_field.filename)
# Move the temporary file to the new filename
new_filename = reqdir + filename
os.rename(tmp_file_path, new_filename)
# check ENV_ID
weights = "/app/weights/best.pt"
device = "cpu"
# support GPU in jetson
if os.getenv("ENV_ID") == "dev":
weights = "/weights/best.pt"
if os.getenv("CUDA_VISIBLE_DEVICES") != "":
device = "0"
run(
weights=weights,
device=device,
source=new_filename,
project=reqdir,
save_txt=True,
save_conf=True,
)
if not os.path.exists(reqdir + "exp/result.txt"):
self.send_response(200) # Send 200 OK status code
self.send_header("Content-type", "application/json")
self.end_headers()
response = {"message": "Nothing found", "result": ""}
self.wfile.write(json.dumps(response).encode("utf-8"))
subprocess.call(["rm", "-rf", reqdir])
return
with open(reqdir + "exp/result.txt", "r") as file:
result = file.read()
response = {"message": "File processed successfully", "result": result}
subprocess.call(["rm", "-rf", reqdir])
self.send_response(200) # Send 200 OK status code
self.send_header("Content-type", "application/json")
self.end_headers()
self.wfile.write(json.dumps(response).encode("utf-8"))
# Create an HTTP server with the request handler
server_address = ("", 8700) # Listen on all available interfaces, port 8700
httpd = ThreadingHTTPServer(server_address, SimpleHTTPRequestHandler)
# Start the server
print("Server running on port 8700")
httpd.serve_forever()