No commit message

XavierJiezou · May 20, 2022 · 4b84af7 · 4b84af7
1 parent 2899f4b
commit 4b84af7
Show file tree

Hide file tree

Showing 8 changed files with 41 additions and 23 deletions.
diff --git a/.flake8 b/.flake8
@@ -16,6 +16,8 @@ extend-ignore =
     INP001
     # ANN101: Missing type annotation for self in method
     ANN101
+    # E402: module level import not at top of file
+    E402
 
 per-file-ignores =
     # F401: Module imported by unused (non-implicit modules)

diff --git a/README.md b/README.md
@@ -1,6 +1,6 @@
 <div align="center">
 
-[![logo](https://raw.githubusercontent.com/XavierJiezou/LitMNIST/main/images/favicon_256x256.svg)](https://pixelied.com/editor/design/6282f5970515730397249959)
+[![logo](https://raw.githubusercontent.com/XavierJiezou/LitMNIST/main/images/logo.png)](https://pixelied.com/editor/design/6282f5970515730397249959)
 
 # LitMNIST
 
@@ -97,7 +97,7 @@ Love the project? Please consider [donating](https://paypal.me/xavierjiezou?coun
 
 ## Demo
 
-![demo](https://cdn.jsdelivr.net/gh/XavierJiezou/LitMNIST@main/images/favicon_256x256.svg)
+![demo](https://cdn.jsdelivr.net/gh/XavierJiezou/LitMNIST@main/images/logo.png)
 
 ## Features
 

diff --git a/images/favicon_280x280.jpg b/images/favicon_280x280.jpg
diff --git a/images/logo.png b/images/logo.png
diff --git a/litmnist/datamodules/mnist_datamodule.py b/litmnist/datamodules/mnist_datamodule.py
@@ -1,4 +1,5 @@
 import os
+from typing import Any
 
 import torch
 from pytorch_lightning import LightningDataModule
@@ -44,9 +45,11 @@ def prepare_data(self) -> None:
         MNIST(self.hparams.data_dir, train=True, download=_download)
         MNIST(self.hparams.data_dir, train=False, download=_download)
 
-    def setup(self, stage: str | None = None) -> None:
+    def setup(self, stage: Any = None) -> None:
         if stage in ("fit", "test", None):
-            trainset = MNIST(root=self.data_dir, train=True, transform=self.transform)
+            trainset = MNIST(
+                root=self.hparams.data_dir, train=True, transform=self.transform
+            )
             testset = MNIST(
                 root=self.hparams.data_dir, train=False, transform=self.transform
             )

diff --git a/litmnist/utils/core.py b/litmnist/utils/core.py
diff --git a/litmnist/utils/make_logo.py b/litmnist/utils/make_logo.py
@@ -0,0 +1,28 @@
+import sys
+
+sys.path.append(".")
+
+from torchvision.utils import save_image
+
+from litmnist.datamodules import MNISTDataModule
+
+
+def make_logo(save_path: str) -> None:
+    datamodule = MNISTDataModule()
+    datamodule.prepare_data()
+    datamodule.setup()
+    temp = {i: [] for i in range(10)}
+    for x, y in datamodule.data_test:
+        if len(temp[y]) < 10:
+            temp[y] += x
+        else:
+            if sum(len(value) for value in temp.values()) == 100:
+                break
+            else:
+                continue
+    logo = [each.unsqueeze(0) for value in temp.values() for each in value]
+    save_image(logo, save_path, nrow=10, padding=0)
+
+
+if __name__ == "__main__":
+    make_logo("images/logo.png")
diff --git a/pyproject.toml b/pyproject.toml
@@ -21,10 +21,13 @@ generate-setup-file = false
 python = "^3.7"
 torch = "^1.8.0"
 torchvision = "^0.9.0"
-pytorch-lightning = "^1.6.3"
+pytorch-lightning = "^1.5.0"
+torchmetrics="^0.8.2"
 comet-ml = "^3.31.0"
 fire = "^0.4.0"
 rich = "^12.3.0"
+hydra-core = "^1.1"
+wandb="^0.12.16"
 
 [tool.poetry.dev-dependencies]
 toml = "^0.10.2"