Merge pull request #17 from lombokai/bugfix/predict

Bugfix/predict

predict.py

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

src/onepiece_classify/infer/recognition.py

@@ -12,8 +12,9 @@
 
 
 class ImageRecognition(BaseInference):
-    def __init__(self, model_path: str):
+    def __init__(self, model_path: str, device: str):
         self.model_path = Path(model_path)
+        self.device = device
         self.class_dict = {
             0: "Ace",
             1: "Akainu",
@@ -39,7 +40,7 @@ def __init__(self, model_path: str):
 
     def _build_model(self):
         # load model
-        state_dict = torch.load(self.model_path)
+        state_dict = torch.load(self.model_path, map_location=self.device)
         model_backbone = image_recog(self.nclass)
         model_backbone.load_state_dict(state_dict)
         return model_backbone
@@ -65,7 +66,7 @@ def pre_process(
         else:
             print("Image type not recognized")
 
-        return img
+        return img.to(self.device)
 
     def forward(self, image_tensor: torch.Tensor) -> torch.Tensor:
         self.model.eval()