display mnist classification chart

display mnist classification chart

pytests

update actions

display mnist classification chart

add fill dir

update license

.github/workflows/lint.yml

@@ -3,21 +3,21 @@ name: black linting
 on:
   push:
     branches:
-      - "main"
-      - "dev"
-      - "test/*"
-      - "test-*"
+      - main
+      - dev
+      - test/*
+      - test-*
   pull_request:
     branches:
-      - "main"
-      - "dev"
+      - main
+      - dev
 
 jobs:
   lint:
     runs-on: ubuntu-latest
     steps:
       - uses: actions/checkout@v3
-      - uses: psf/black@latest
+      - uses: psf/black@stable
         with:
           options: "--check --verbose"
           src: "."

.github/workflows/test.yml

@@ -3,14 +3,17 @@ name: test repository
 on:
   push:
     branches:
-      - "main"
-      - "dev"
-      - "test/*"
-      - "test-*"
+      - main
+      - dev
+      - test/*
+      - test-*
   pull_request:
     branches:
-      - "main"
-      - "dev"
+      - main
+      - dev
+
+permissions:
+  contents: read
 
 jobs:
   test:

LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2024 Arnold Unterauer
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

README.md

@@ -1 +1,7 @@
-# MNIST
+# Metis - Machine Learning Explained through Interactive Simulations
+
+![MIT license](https://img.shields.io/badge/license-MIT-blue)
+[![Code style: black](https://img.shields.io/badge/code%20style-black-000000.svg)](https://github.com/psf/black)
+
+[![Black Linting](https://github.com/Anemosx/metis/actions/workflows/lint.yml/badge.svg)](https://github.com/Anemosx/metis/actions/workflows/lint.yml)
+[![PyTest](https://github.com/Anemosx/metis/actions/workflows/test.yml/badge.svg)](https://github.com/Anemosx/metis/actions/workflows/test.yml)

app/routers/predict.py

@@ -44,9 +44,16 @@ async def predict_digit(data: ImageData) -> JSONResponse:
         image_bytes = base64.b64decode(data.image.split(",")[1])
         image = Image.open(io.BytesIO(image_bytes))
 
+        if image.mode != "RGBA":
+            image = image.convert("RGBA")
+
+        image = image.getchannel("A").resize((28, 28), Image.NEAREST)
+
         predicted_class, prediction_dist = predict(image)
 
-        return JSONResponse(content={"prediction": predicted_class})
+        return JSONResponse(
+            content={"prediction": predicted_class, "distribution": prediction_dist}
+        )
 
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e))

app/vision/mnist.py

@@ -1,9 +1,11 @@
+import io
 import os
 
+import numpy as np
 import torch
 from PIL.Image import Image
 from torch import nn
-from torchvision import transforms
+import torch.nn.functional as F
 
 mnist_data = {}
 
@@ -65,7 +67,11 @@ def forward(self, inputs: torch.Tensor) -> torch.Tensor:
         return out
 
 
-def load_mnist_model(model: nn.Module, filename: str | None = None) -> None:
+def load_mnist_model(
+    model: nn.Module,
+    filename: str | None = None,
+    device: str | torch.device | None = None,
+) -> None:
     """
     Load MNIST model weights from a specified file into the provided model instance.
 
@@ -75,12 +81,17 @@ def load_mnist_model(model: nn.Module, filename: str | None = None) -> None:
         The PyTorch model instance into which the weights will be loaded.
     filename : str, optional
         The path to the file containing the model weights.
+    device : str, torch.device or None, optional
+        Device to run the evaluation on.
     """
 
+    if device is None:
+        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
     if filename is None:
         filename = os.path.join(os.getcwd(), "models", "mnist", "model.pt")
 
-    checkpoint = torch.load(filename)
+    checkpoint = torch.load(filename, map_location=device)
     model.load_state_dict(checkpoint["model_state_dict"])
 
 
@@ -95,17 +106,8 @@ async def init_vision_model() -> None:
     model.eval()
     load_mnist_model(model)
 
-    mnist_transform = transforms.Compose(
-        [
-            transforms.Resize((28, 28)),
-            transforms.ToTensor(),
-            transforms.Normalize((0.5,), (0.5,)),
-        ]
-    )
-
     mnist_data["device"] = device
     mnist_data["model"] = model
-    mnist_data["transform"] = mnist_transform
 
 
 def predict(image: Image) -> tuple[int, list[float]]:
@@ -125,15 +127,16 @@ def predict(image: Image) -> tuple[int, list[float]]:
         The raw output logits for each digit class.
     """
 
+    alpha_np = np.array(image, dtype=np.float32) / 255.0
+
     image_tensor = (
-        mnist_data["transform"](image)
-        .unsqueeze(0)
-        .to(mnist_data["device"])[:, [3], :, :]
+        torch.tensor(alpha_np).unsqueeze(0).unsqueeze(0).to(mnist_data["device"])
     )
+    image_tensor = (image_tensor - 0.5) * 2
 
     with torch.no_grad():
         output = mnist_data["model"](image_tensor)[0]
-        output_list = output.to("cpu").tolist()
+        output_list = F.softmax(output, dim=0).to("cpu").tolist()
         predicted_class = torch.argmax(output).item()
 
     return predicted_class, output_list