Skip to content
Projects
Groups
Snippets
Help
Loading...
Help
Support
Keyboard shortcuts
?
Submit feedback
Contribute to GitLab
Sign in / Register
Toggle navigation
2022-154
Project overview
Project overview
Details
Activity
Releases
Repository
Repository
Files
Commits
Branches
Tags
Contributors
Graph
Compare
Issues
0
Issues
0
List
Boards
Labels
Milestones
Merge Requests
1
Merge Requests
1
CI / CD
CI / CD
Pipelines
Jobs
Schedules
Analytics
Analytics
CI / CD
Repository
Value Stream
Wiki
Wiki
Snippets
Snippets
Members
Members
Collapse sidebar
Close sidebar
Activity
Graph
Create a new issue
Jobs
Commits
Issue Boards
Open sidebar
2022-154
2022-154
Commits
a9b3f7e9
Commit
a9b3f7e9
authored
Oct 09, 2022
by
Senatilaka T.S.
Browse files
Options
Browse Files
Download
Email Patches
Plain Diff
Upload New File
parent
31628f88
Changes
1
Hide whitespace changes
Inline
Side-by-side
Showing
1 changed file
with
111 additions
and
0 deletions
+111
-0
main.py
main.py
+111
-0
No files found.
main.py
0 → 100644
View file @
a9b3f7e9
import
io
import
json
import
base64
import
uuid
import
uvicorn
from
typing
import
List
from
PIL
import
Image
from
io
import
BytesIO
from
starlette.responses
import
Response
from
fastapi
import
FastAPI
,
File
,
WebSocket
,
WebSocketDisconnect
from
yolov5
import
yolov5
,
get_image_from_bytes
# FastAPI
app
=
FastAPI
(
title
=
"Strawberry Diseases Model"
,
description
=
"""Visit port 8088/docs for the FastAPI documentation."""
,
version
=
"0.0.1"
,
)
class
ConnectionManager
:
"""Web socket connection manager."""
def
__init__
(
self
):
self
.
active_connections
:
List
[
WebSocket
]
=
[]
async
def
connect
(
self
,
websocket
:
WebSocket
):
await
websocket
.
accept
()
self
.
active_connections
.
append
(
websocket
)
def
disconnect
(
self
,
websocket
:
WebSocket
):
self
.
active_connections
.
remove
(
websocket
)
async
def
send_message
(
self
,
message
:
str
,
websocket
:
WebSocket
):
await
websocket
.
send_text
(
message
)
async
def
broadcast
(
self
,
message
:
str
):
for
connection
in
self
.
active_connections
:
await
connection
.
send_text
(
message
)
conn_mgr
=
ConnectionManager
()
def
base64_encode_img
(
img
):
buffered
=
BytesIO
()
img
.
save
(
buffered
,
format
=
"PNG"
)
buffered
.
seek
(
0
)
img_byte
=
buffered
.
getvalue
()
encoded_img
=
"data:image/png;base64,"
+
base64
.
b64encode
(
img_byte
)
.
decode
()
return
encoded_img
@
app
.
get
(
"/"
)
def
home
():
return
{
"message"
:
"YOLO - You Only Look Once"
}
@
app
.
post
(
"/yolo"
)
def
process_yolov5
(
file
:
bytes
=
File
(
...
)):
# file_bytes = file.file.read()
# image = Image.open(io.BytesIO(file_bytes))
# name = "/data/temp.png"
# # image.save(name)
# image.filename = name
classes
,
converted_img
=
yolov5
(
get_image_from_bytes
(
file
))
bytes_io
=
io
.
BytesIO
()
converted_img
.
save
(
'test.png'
)
converted_img
.
save
(
bytes_io
,
format
=
"PNG"
)
rHeaders
=
{
"Detected Classes"
:
classes
[
0
]}
return
Response
(
content
=
bytes_io
.
getvalue
(),
headers
=
rHeaders
,
media_type
=
"image/png"
)
@
app
.
websocket
(
"/yolo_ws/{client_id}"
)
async
def
process_yolov5_ws
(
websocket
:
WebSocket
,
client_id
:
int
):
await
conn_mgr
.
connect
(
websocket
)
try
:
while
True
:
data
=
await
websocket
.
receive_text
()
# Convert to PIL image
image
=
data
[
data
.
find
(
","
)
+
1
:]
dec
=
base64
.
b64decode
(
image
+
"==="
)
image
=
Image
.
open
(
BytesIO
(
dec
))
.
convert
(
"RGB"
)
# Process the image
name
=
f
"/data/{str(uuid.uuid4())}.png"
image
.
filename
=
name
classes
,
converted_img
=
yolov5
(
image
)
result
=
{
"prediction"
:
json
.
dumps
(
classes
),
"output"
:
base64_encode_img
(
converted_img
),
}
# logging.info("-----", json.dumps(result))
# Send back the result
await
conn_mgr
.
send_message
(
json
.
dumps
(
result
),
websocket
)
# await conn_mgr.broadcast(f"Client #{client_id} says: {data}")
except
WebSocketDisconnect
:
conn_mgr
.
disconnect
(
websocket
)
await
conn_mgr
.
broadcast
(
f
"Client #{client_id} left the chat"
)
if
__name__
==
"__main__"
:
uvicorn
.
run
(
"main:app"
,
host
=
"0.0.0.0"
,
port
=
8088
,
reload
=
True
)
Write
Preview
Markdown
is supported
0%
Try again
or
attach a new file
Attach a file
Cancel
You are about to add
0
people
to the discussion. Proceed with caution.
Finish editing this message first!
Cancel
Please
register
or
sign in
to comment