Merge pull request #9 from lombokai/feature/inference

Feature/inference

Author: nunenuh

.gitignore

@@ -1,12 +1,13 @@
 # ignore dataset
 raw_data/*
-src/data/train/*
-src/data/val/*
+data/train/*
+data/val/*
 
 # Byte-compiled / optimized / DLL files
 __pycache__/
 *.py[cod]
 *$py.class
+.vscode
 
 # C extensions
 *.so
@@ -125,6 +126,7 @@ celerybeat.pid
 *.sage.py
 
 # Environments
+.myenv
 .env
 .venv
 env/
@@ -164,4 +166,4 @@ cython_debug/
 #  option (not recommended) you can uncomment the following to ignore the entire idea folder.
 #.idea/
 
-data/
+# data/

src/onepiece_classify/data/data_setup.py

@@ -7,6 +7,7 @@
 )
 from pathlib import Path
 
+
 class OnepieceImageDataLoader:
     def __init__(
         self,
@@ -67,6 +68,3 @@ def _build_dataloader(self, mode: str = "train", shuffle: bool = True) -> DataLo
         )
 
         return loader
-
-
-

src/onepiece_classify/data/setup_data.py

@@ -1 +1,2 @@
-from .base import *
+from .base import *
+from .recognition import *

src/onepiece_classify/infer/__init__.py

@@ -8,7 +8,7 @@
 class BaseInference(ABC):
 
     @abstractmethod
-    def pre_process(self, image: Optional[str, np.ndarray, Image])->torch.Tensor:
+    def pre_process(self, image: Optional[str|np.ndarray|Image.Image])->torch.Tensor:
         pass
 
     @abstractmethod
@@ -20,5 +20,5 @@ def post_process(self, output: torch.Tensor) -> str:
         pass
 
     @abstractmethod
-    def predict(self, image: Optional[str, np.ndarray, Image]) -> dict:
-        pass
+    def predict(self, image: Optional[str|np.ndarray|Image.Image]) -> dict:
+        pass

src/onepiece_classify/infer/base.py

@@ -1,28 +1,92 @@
+from PIL import Image
+import numpy as np
+import torch
+from pathlib import Path
+from .base import BaseInference
+from typing import Optional, Tuple, Dict
 
+from onepiece_classify.models import image_recog
+from onepiece_classify.transforms import get_test_transforms
 
 
-
-from .base import BaseInference
-import torch
-
 class ImageRecognition(BaseInference):
 
     def __init__(self, model_path: str):
-        self.model_path = model_path
+        self.model_path = Path(model_path)
+        self.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',
+        }
+        self.nclass = len(self.class_dict)
+        self.model = self._build_model()
 
-        self.model = self._build_model(model_path)
-        
-    def _build_model(self, model_path):
+    def _build_model(self):
          # load model
-        state_dict = torch.load(model_path)
-
-        model = ImageModel.load(state_dict)
-        return model
+        state_dict = torch.load(self.model_path)
+        model_backbone = image_recog(self.nclass)
+        model_backbone.load_state_dict(state_dict)
+        return model_backbone
 
+    def pre_process(self, image: Optional[str | np.ndarray | Image.Image]) -> torch.Tensor:
+        
+        trans = get_test_transforms()
+
+        if isinstance(image, str):
+            img = Image.open(image).convert("RGB")
+            img = trans(img).unsqueeze(0)
+        
+        elif isinstance(image, Image.Image):
+            img = image.convert("RGB")
+            img = trans(img).unsqueeze(0)
+
+        elif isinstance(image, np.ndarray):
+            img = image.astype(np.uint8)
+            img = Image.fromarray(img).convert("RGB")
+            img = trans(img).unsqueeze(0)
+
+        else:
+            print("Image type not recognized")
+
+        return img
+
     def forward(self, image_tensor: torch.Tensor) -> torch.Tensor:
-        return self.model.forward(image_tensor)
-    
+        self.model.eval()
+
+        result = self.model(image_tensor)
+        return result
 
+    def post_process(self, output: torch.Tensor) -> Tuple[str, float]:
+        
+        logits_prob = torch.softmax(output, dim=1).squeeze()
+        class_idx = int(torch.argmax(logits_prob))
+        
+        class_names = self.class_dict[class_idx]
+        confidence = logits_prob[class_idx]
+        return (class_names, float(confidence))
 
-def recognition():
-    return ImageRecognition(model_path="src/checkpoint/checkpoint_notebook.pth")
+    def predict(self, image: Optional[str|np.ndarray|Image.Image]) -> Dict[str, str]:
+        
+        tensor_img = self.pre_process(image=image)
+        logits = self.forward(tensor_img)
+        class_names, confidence = self.post_process(logits)
+
+        return {
+            "class_names": class_names,
+            "confidence": f"{confidence:.4f}"
+        }

src/onepiece_classify/infer/inference.py

@@ -7,29 +7,27 @@ class ImageRecogModel(nn.Module):
 
     def __init__(self, num_classes):
         super().__init__()
-
         self.num_classes = num_classes
-        self.backbone = self._build_backbone()
-        self.in_features = self._build_backbone().classifier[0].in_features
-        self.backbone.classifier = nn.Sequential(
-            nn.Dropout(p=0.2),
-            nn.Linear(self.in_features, out_features=self.num_classes)
-        )
-
-        # self.dropout = nn.Dropout(0.2)
 
-    def _build_backbone(self):
+    def build_backbone(self):
         model = models.mobilenet_v3_large(weights="DEFAULT")
 
         for param in model.parameters():
             param.requires_grad = False
 
+        in_features = model.classifier[0].in_features
+        model.classifier = nn.Sequential(
+            nn.Dropout(p=0.2),
+            nn.Linear(in_features, out_features=self.num_classes)
+        )
+
         return model
 
     def forward(self, x):
-        x = self.backbone(x)        
+        model = self.build_backbone()
+        x = model(x)        
         return x
 
 def image_recog(num_classes):
-    net = ImageRecogModel(num_classes)
+    net = ImageRecogModel(num_classes).build_backbone()
     return net