added semaphore for cnn server
This commit is contained in:
Binary file not shown.
+16
-11
@@ -11,9 +11,13 @@ import torch.nn.functional as F
|
|||||||
|
|
||||||
from bird_cnn import Bird_CNN
|
from bird_cnn import Bird_CNN
|
||||||
|
|
||||||
|
import threading
|
||||||
|
|
||||||
BUILD_PATH = "./build_models"
|
BUILD_PATH = "./build_models"
|
||||||
IMAGE_SIZE = 256
|
IMAGE_SIZE = 256
|
||||||
|
|
||||||
|
sem = threading.Semaphore(1) #adjust to performance
|
||||||
|
|
||||||
class bird_species(Enum):
|
class bird_species(Enum):
|
||||||
Common_Kingfisher = 0
|
Common_Kingfisher = 0
|
||||||
CommonMyna = 1
|
CommonMyna = 1
|
||||||
@@ -53,17 +57,18 @@ app.add_middleware(
|
|||||||
|
|
||||||
@app.post("/predict")
|
@app.post("/predict")
|
||||||
async def predict(file: UploadFile = File(...)):
|
async def predict(file: UploadFile = File(...)):
|
||||||
image_bytes = await file.read()
|
with sem:
|
||||||
|
image_bytes = await file.read()
|
||||||
|
|
||||||
image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
|
image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
|
||||||
image = transform(image).unsqueeze(0).to(device)
|
image = transform(image).unsqueeze(0).to(device)
|
||||||
|
|
||||||
with torch.no_grad():
|
with torch.no_grad():
|
||||||
pred = model(image)
|
pred = model(image)
|
||||||
probs = F.softmax(pred, dim=1)
|
probs = F.softmax(pred, dim=1)
|
||||||
confidence, cls = torch.max(probs, dim=1)
|
confidence, cls = torch.max(probs, dim=1)
|
||||||
|
|
||||||
return {
|
return {
|
||||||
"class": bird_species(cls.item()).name,
|
"class": bird_species(cls.item()).name,
|
||||||
"confidence": confidence.item()
|
"confidence": confidence.item()
|
||||||
}
|
}
|
||||||
Reference in New Issue
Block a user