fix building
This commit is contained in:
parent
8c296400f2
commit
1a748bcc62
2
Makefile
2
Makefile
|
|
@ -1,5 +1,5 @@
|
||||||
api:
|
api:
|
||||||
uv run granian --interface asgi server.main:app
|
uv run granian --interface asgi server.main:app --host 0.0.0.0
|
||||||
|
|
||||||
dog-train:
|
dog-train:
|
||||||
uv run ml/dogs.py
|
uv run ml/dogs.py
|
||||||
|
|
|
||||||
|
|
@ -2,9 +2,7 @@ version: "3"
|
||||||
|
|
||||||
services:
|
services:
|
||||||
server:
|
server:
|
||||||
image: aibeerds:latest
|
image: ai:latest
|
||||||
ports:
|
ports:
|
||||||
- 8000:8000
|
- 8000:8000
|
||||||
volumes:
|
|
||||||
- ./:/app
|
|
||||||
restart: always
|
restart: always
|
||||||
|
|
@ -35,7 +35,6 @@ def load_model(model_path, device="cpu"):
|
||||||
with open("server/meta/images.json", "r") as f:
|
with open("server/meta/images.json", "r") as f:
|
||||||
IMAGES = json.loads(f.read())
|
IMAGES = json.loads(f.read())
|
||||||
|
|
||||||
|
|
||||||
MODEL = load_model("server/models/dogs_model.pth")
|
MODEL = load_model("server/models/dogs_model.pth")
|
||||||
|
|
||||||
with open("server/meta/labels_dogs.json", "r") as f:
|
with open("server/meta/labels_dogs.json", "r") as f:
|
||||||
|
|
@ -43,7 +42,7 @@ with open("server/meta/labels_dogs.json", "r") as f:
|
||||||
labels_dict = json.loads(data_labels)
|
labels_dict = json.loads(data_labels)
|
||||||
|
|
||||||
|
|
||||||
def predict_image(image, model, device="cuda"):
|
def predict_image(image, model, device="cuda") -> list[tuple]:
|
||||||
img_size = (180, 180)
|
img_size = (180, 180)
|
||||||
preprocess = transforms.Compose(
|
preprocess = transforms.Compose(
|
||||||
[
|
[
|
||||||
|
|
@ -57,11 +56,15 @@ def predict_image(image, model, device="cuda"):
|
||||||
|
|
||||||
with torch.no_grad():
|
with torch.no_grad():
|
||||||
output = model(input_batch)
|
output = model(input_batch)
|
||||||
|
data = []
|
||||||
|
for i in range(3):
|
||||||
|
if i >= len(output):
|
||||||
|
break
|
||||||
|
probabilities = F.softmax(output[i], dim=0)
|
||||||
|
|
||||||
probabilities = F.softmax(output[0], dim=0)
|
_, predicted_idx = torch.max(probabilities, 0)
|
||||||
|
data.append((predicted_idx.item(), probabilities.cpu().numpy()))
|
||||||
_, predicted_idx = torch.max(probabilities, 0)
|
return data
|
||||||
return predicted_idx.item(), probabilities.cpu().numpy()
|
|
||||||
|
|
||||||
|
|
||||||
class BeerdsController(Controller):
|
class BeerdsController(Controller):
|
||||||
|
|
@ -74,12 +77,17 @@ class BeerdsController(Controller):
|
||||||
body = await data.read()
|
body = await data.read()
|
||||||
|
|
||||||
img_file = Image.open(io.BytesIO(body))
|
img_file = Image.open(io.BytesIO(body))
|
||||||
predicted_idx, probabilities = predict_image(img_file, MODEL, "cpu")
|
predicted_data = predict_image(img_file, MODEL, "cpu")
|
||||||
predicted_label = labels_dict[str(predicted_idx)]
|
results = {}
|
||||||
|
images = []
|
||||||
images = [{"name": predicted_label, "url": IMAGES[predicted_label]}]
|
for d in predicted_data:
|
||||||
|
predicted_idx, probabilities = d
|
||||||
|
predicted_label = labels_dict[str(predicted_idx)]
|
||||||
|
name = predicted_label.replace("_", " ")
|
||||||
|
images.append({"name": name, "url": IMAGES[name]})
|
||||||
|
results[float(probabilities[0])] = name
|
||||||
return {
|
return {
|
||||||
"results": {float(probabilities[0]): predicted_label},
|
"results": results,
|
||||||
"images": images,
|
"images": images,
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
|
||||||
Loading…
Reference in New Issue