get several results + deploy

.gitea/workflows/ai.yaml

@@ -21,4 +21,4 @@ jobs:
         with:
           context: .
           push: true
-          tags: https://gitea.webart-tech.ru/webart/beerds:${{ gitea.sha }}
+          tags: https://gitea.webart-tech.ru/webart/beerds/backend:${{ gitea.sha }}

docker-compose.yml

@@ -5,4 +5,4 @@ services:
     image: ai:latest
     ports:
       - 8000:8000
-    restart: always
+    restart: 'no'

ml/dogs.py

@@ -12,7 +12,7 @@ ImageFile.LOAD_TRUNCATED_IMAGES = True
 print(f"Using device: {DEVICE}")
 IMG_SIZE = (180, 180)
 INPUT_DIR = "assets/dog"
-NUM_EPOCHS = 10
+NUM_EPOCHS = 100
 MODEL_NAME = "dogs_model.pth"
 
 if __name__ == "__main__":

server/main.py

@@ -61,14 +61,14 @@ def predict_image(image, model, device="cuda") -> list[tuple]:
 
     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)
 
-        _, predicted_idx = torch.max(probabilities, 0)
-        data.append((predicted_idx.item(), probabilities.cpu().numpy()))
+    probabilities = torch.nn.functional.softmax(output[0], dim=0)
+    k = 5
+    topk_probs, predicted_idx = torch.topk(probabilities, k)
+
+    data = []
+    for i in range(k):
+        data.append((predicted_idx[i].item(), float(topk_probs[i].item())))
     return data
 
 
@@ -90,7 +90,7 @@ class BeerdsController(Controller):
             predicted_label = labels_dogs[str(predicted_idx)]
             name = predicted_label.replace("_", " ")
             images.append({"name": name, "url": IMAGES[name]})
-            results[float(probabilities[0])] = name
+            results[probabilities] = name
         return {
             "results": results,
             "images": images,
@@ -108,7 +108,7 @@ class BeerdsController(Controller):
         for d in predicted_data:
             predicted_idx, probabilities = d
             predicted_label = labels_cats[str(predicted_idx)]
-            results[float(probabilities[0])] = predicted_label
+            results[probabilities] = predicted_label
         return {
             "results": results,
             "images": [],

server/static/scripts.js

@@ -13,7 +13,7 @@ async function SavePhoto(self) {
     let text = "<h3 class='image-results'>Результаты</h3>";
     let uniqChecker = {};
     for (let key in json.results) {
-      text += "<div class='image-block'><div class='image-text'>" + json.results[key] + " (вероятность: " +  Math.round(parseFloat(key)*10)/10 + ")</div>";
+      text += "<div class='image-block'><div class='image-text'>" + json.results[key] + " (вероятность: " +  Math.round(parseFloat(key)*100) + "%)</div>";
       for (let imgKey in json.images) {
         if (json.images[imgKey].name != json.results[key]) {
           continue;