gallery

.gitattributes

@@ -1 +1,2 @@
 *.pth filter=lfs diff=lfs merge=lfs -text
+*.webp filter=lfs diff=lfs merge=lfs -text

server/main.py

@@ -32,9 +32,6 @@ def load_model(model_path, device="cpu"):
     return model
 
 
-with open("server/meta/images.json", "r") as f:
-    IMAGES = json.loads(f.read())
-
 DOG_MODEL = load_model("server/models/dogs_model.pth")
 CAT_MODEL = load_model("server/models/cats_model.pth")
 
@@ -47,6 +44,9 @@ with open("server/meta/labels_cats.json", "r") as f:
     labels_cats = json.loads(data_labels)
 
 
+with open("server/meta/images.json", "r") as f:
+    IMAGES = json.loads(f.read())
+
 def predict_image(image, model, device="cuda") -> list[tuple]:
     img_size = (224, 224)
     preprocess = transforms.Compose(
@@ -89,7 +89,10 @@ class BeerdsController(Controller):
             predicted_idx, probabilities = d
             predicted_label = labels_dogs[str(predicted_idx)]
             name = predicted_label.replace("_", " ")
-            images.append({"name": name, "url": IMAGES[name]})
+            images.append({
+                "name": name, 
+                "url": [f"/static/assets/dog/{predicted_label}/{i}" for i in IMAGES['dog'][predicted_label]]
+            })
             results[probabilities] = name
         return {
             "results": results,
@@ -105,13 +108,19 @@ class BeerdsController(Controller):
         img_file = Image.open(io.BytesIO(body))
         predicted_data = predict_image(img_file, CAT_MODEL, "cpu")
         results = {}
+        images = []
         for d in predicted_data:
             predicted_idx, probabilities = d
             predicted_label = labels_cats[str(predicted_idx)]
+            name = predicted_label.replace("_", " ")
+            images.append({
+                "name": name, 
+                "url": [f"/static/assets/cat/{predicted_label}/{i}" for i in IMAGES['cat'][predicted_label]]
+            })
             results[probabilities] = predicted_label
         return {
             "results": results,
-            "images": [],
+            "images": images,
         }