zjt/SAMapi
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Compare commits

base: zjt:de39a53ee043be4b5ee70c3197456898391a48d2
Branches Tags
zjt:main
..
compare: zjt:1b4cec53f56156a8cd4648a37f8788806848b600
Branches Tags
zjt:main

2 Commits
de39a53ee0 ... 1b4cec53f5

Author SHA1 Message Date
zjt
1b4cec53f5 cros 2023-12-18 10:06:41 +08:00
zjt
2aca3657ad fix 2023-12-18 10:05:33 +08:00
1 changed files with 8 additions and 4 deletions
Show Stats Expand all files Collapse all files

12
backend.py
View File

@ -13,7 +13,7 @@ print(torch.cuda.is_available())
# checkpoint = "sam_vit_b_01ec64.pth" # checkpoint = "sam_vit_b_01ec64.pth"
print(time.time()) print(time.time())
# set large model related configs # set large model related configs
checkpoint = "sam_vit_h_4b8939.pth" checkpoint = "notebooks/sam_vit_h_4b8939.pth"
model_type = "vit_h" model_type = "vit_h"
sam = sam_model_registry[model_type](checkpoint=checkpoint) sam = sam_model_registry[model_type](checkpoint=checkpoint)
sam.to(device='cuda') sam.to(device='cuda')
@ -51,9 +51,13 @@ def index():
np.save("array.npy", image_embedding) np.save("array.npy", image_embedding)
response = Response(byte_io, mimetype="application/octet-stream") response = Response(byte_io, mimetype="application/octet-stream")
response.headers["Content-Length"] = 4194432 response.headers["Content-Length"] = 4194432
response.headers["Access-Control-Allow-Origin"] = "*"
return response return response
else: else:
return 'Content-Type not supported!' return 'Content-Type not supported!'
print('server starts working') print('server starts working')
app.run('0.0.0.0', 2020, False)