diff --git a/Makefile b/Makefile index 36812dd..09a4f19 100644 --- a/Makefile +++ b/Makefile @@ -1,5 +1,5 @@ api: - uv run granian --interface asgi server.main:app + uv run granian --interface asgi server.main:app --host 0.0.0.0 dog-train: uv run ml/dogs.py diff --git a/docker-compose.yml b/docker-compose.yml index efb5753..a65f8ea 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -2,9 +2,7 @@ version: "3" services: server: - image: aibeerds:latest + image: ai:latest ports: - 8000:8000 - volumes: - - ./:/app restart: always \ No newline at end of file diff --git a/server/main.py b/server/main.py index 6062148..2fca7ff 100644 --- a/server/main.py +++ b/server/main.py @@ -35,7 +35,6 @@ def load_model(model_path, device="cpu"): with open("server/meta/images.json", "r") as f: IMAGES = json.loads(f.read()) - MODEL = load_model("server/models/dogs_model.pth") 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) -def predict_image(image, model, device="cuda"): +def predict_image(image, model, device="cuda") -> list[tuple]: img_size = (180, 180) preprocess = transforms.Compose( [ @@ -57,11 +56,15 @@ def predict_image(image, model, device="cuda"): with torch.no_grad(): 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) - return predicted_idx.item(), probabilities.cpu().numpy() + _, predicted_idx = torch.max(probabilities, 0) + data.append((predicted_idx.item(), probabilities.cpu().numpy())) + return data class BeerdsController(Controller): @@ -74,12 +77,17 @@ class BeerdsController(Controller): body = await data.read() img_file = Image.open(io.BytesIO(body)) - predicted_idx, probabilities = predict_image(img_file, MODEL, "cpu") - predicted_label = labels_dict[str(predicted_idx)] - - images = [{"name": predicted_label, "url": IMAGES[predicted_label]}] + predicted_data = predict_image(img_file, MODEL, "cpu") + results = {} + images = [] + 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 { - "results": {float(probabilities[0]): predicted_label}, + "results": results, "images": images, }