From 486e5d0d57dd117c5674c18a8d5184bb13c6853a Mon Sep 17 00:00:00 2001
From: artem <artem@webart-tech.ru>
Date: Wed, 23 Apr 2025 18:51:17 +0300
Subject: [PATCH] get several results + deploy

---
 .gitea/workflows/ai.yaml |  2 +-
 docker-compose.yml       |  2 +-
 ml/dogs.py               |  2 +-
 server/main.py           | 18 +++++++++---------
 server/static/scripts.js |  2 +-
 5 files changed, 13 insertions(+), 13 deletions(-)

diff --git a/.gitea/workflows/ai.yaml b/.gitea/workflows/ai.yaml
index b24d714..8fea5d4 100644
--- a/.gitea/workflows/ai.yaml
+++ b/.gitea/workflows/ai.yaml
@@ -21,4 +21,4 @@ jobs:
         with:
           context: .
           push: true
-          tags: https://gitea.webart-tech.ru/webart/beerds:${{ gitea.sha }}
\ No newline at end of file
+          tags: https://gitea.webart-tech.ru/webart/beerds/backend:${{ gitea.sha }}
\ No newline at end of file
diff --git a/docker-compose.yml b/docker-compose.yml
index a65f8ea..b141711 100644
--- a/docker-compose.yml
+++ b/docker-compose.yml
@@ -5,4 +5,4 @@ services:
     image: ai:latest
     ports:
       - 8000:8000
-    restart: always
\ No newline at end of file
+    restart: 'no'
\ No newline at end of file
diff --git a/ml/dogs.py b/ml/dogs.py
index adc1a99..6d561aa 100644
--- a/ml/dogs.py
+++ b/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__":
diff --git a/server/main.py b/server/main.py
index a91d8cc..a905ac2 100644
--- a/server/main.py
+++ b/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": [],
diff --git a/server/static/scripts.js b/server/static/scripts.js
index a68d47a..db2c2a1 100644
--- a/server/static/scripts.js
+++ b/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;