fix building

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

docker-compose.yml

@@ -2,9 +2,7 @@ version: "3"
 
 services:
   server:
-    image: aibeerds:latest
+    image: ai:latest
     ports:
       - 8000:8000
-    volumes:
-      - ./:/app
     restart: always

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,
         }