Feat/train

predict.py

@@ -12,7 +12,7 @@ def predict_image(
     device: str = typer.Argument("cpu", help="use cuda if your device has cuda", show_default=True)
 ):
 
-    predictor = ImageRecognition(model_path=model_path, manifest_path="checkpoint/manifest.pth", device=device)
+    predictor = ImageRecognition(model_path=model_path, device=device)
     result = predictor.predict(image=image_path)
     typer.echo(f"Prediction: {result}")
 

src/lightning_classify/infer/infer.py

@@ -8,18 +8,20 @@
 import torch.nn as nn
 from PIL import Image
 
-from lightning_classify.utils import downloader
+from lightning_classify.utils import downloader, unzip
 from lightning_classify.models import ImageRecogModelV1
 from lightning_classify.transforms import ImageTransform
 
 
 class ImageRecognition:
-    def __init__(self, device: torch.device, model_path=None, manifest_path=None):
-        
-        path_to_save = str(self._get_cache_dir()) + "/model.ckpt"
+    def __init__(self, device: torch.device, model_path=None):
+
+        path_to_save = str(self._get_cache_dir()) + "/cache.zip"
         if model_path is None:
             downloader(path_to_save)
-            self.model_path = path_to_save
+            unzip(path_to_save, to_save=str(self._get_cache_dir()))
+            self.model_path = str(self._get_cache_dir()) + "/model.ckpt"
+            manifest_path = str(self._get_cache_dir()) + "/manifest.pth"
         else:
             self.model_path = Path(model_path)
             if not self.model_path.exists():
@@ -30,7 +32,7 @@ def __init__(self, device: torch.device, model_path=None, manifest_path=None):
                                                             criterion=nn.CrossEntropyLoss(),
                                                             map_location=self.device)
 
-        self.class_dict = torch.load(manifest_path)
+        self.class_dict = torch.load(manifest_path)["class_dict"]
 
     def _get_cache_dir(self):
         if sys.platform.startswith("win"):

src/lightning_classify/utils/__init__.py

@@ -1 +1,2 @@
-from .downloader import *
+from .downloader import *
+from .unzip_file import *

src/lightning_classify/utils/downloader.py

@@ -4,7 +4,7 @@
 
 def downloader(output_file):
 
-    file_id = "1XqKMk28A-gvitoztaeH20X-dOqLNtraT" # not ignore criterion
+    file_id = "1JGUYPu5xrAg736ViylLAXi_TCro_bjQk"
     prefix = 'https://drive.google.com/uc?/export=download&id='
 
     url_download = prefix+file_id

src/lightning_classify/utils/unzip_file.py

@@ -0,0 +1,9 @@
+import zipfile
+
+
+def unzip(zip_file_path, to_save):
+
+    with zipfile.ZipFile(zip_file_path, 'r') as zip_ref:
+        zip_ref.extractall(to_save)
+
+

train.py

@@ -24,6 +24,7 @@
     batch_size=32,
     num_workers=11
 )
+loader.setup()
 
 model = ImageRecogModelV1(
     lrate = args.learning_rate,
@@ -35,27 +36,10 @@
                     accelerator=args.accelerator)
 
 # dict to save
-class_dict = {
-    0 : "Ace",
-    1 : "Akainu",
-    2 : "Brook",
-    3 : "Chopper",
-    4 : "Crocodile",
-    5 : "Franky",
-    6 : "Jinbei",
-    7 : "Kurohige",
-    8 : "Law",
-    9 : "Luffy",
-    10 : "Mihawk",
-    11 : "Nami",
-    12 : "Rayleigh",
-    13 : "Robin",
-    14 : "Sanji",
-    15 : "Shanks",
-    16 : "Usopp",
-    17 : "Zoro"
+manifest_dict = {
+    "class_dict": {i:val for (val, i) in loader.class_to_idx.items()}
 }
 
 if __name__ == "__main__":
-    trainer.fit(model, loader)
-    torch.save(class_dict, "checkpoint/manifest.pth", )
+    # trainer.fit(model, loader)
+    torch.save(manifest_dict, "checkpoint/manifest.pth", )