Migrate Image Training Endpoint to Django

training/training/core/dataset.py

@@ -12,6 +12,9 @@
 import torch
 from torch.utils.data import Dataset
 from torch.autograd import Variable
+from torchvision import datasets, transforms
+from torch.util.data import DataLoader, random_split
+from enum import Enum
 
 
 class TrainTestDatasetCreator(ABC):
@@ -98,3 +101,39 @@ def getCategoryList(self) -> list[str]:
         if self._category_list is None:
             raise Exception("Category list not available")
         return self._category_list
+
+class DefaultImageDatasets(Enum):
+    MNIST = "MNIST"
+    FASHION_MNIST = "FashionMNIST"
+    KMNIST = "KMNIST"
+class ImageDefaultDatasetCreator(TrainTestDatasetCreator):
+
+    def __init__(
+        self, 
+        dataset_name:str, 
+        train_tarnsform:None,
+        test_transform:None,
+        batch_size: int = 32, 
+        shuffle: bool = True,
+    ):
+        if dataset_name not in DefaultImageDatasets.__members__:
+            raise Exception(f"The {dataset_name} file does not currently exist in our inventory. Please submit a request to the contributors of the repository")
+        self.train_transform = train_tarnsform or transforms.Compose([transforms.toTensor()])
+        self.test_transform = test_transform or transforms.Compose([transforms.toTensor()])
+        self.batch_size = batch_size
+        self.shuffle = shuffle
+
+        self.train_set = datasets.__dict__[dataset_name](root='./backend/image_data_uploads', train=True, download=True, transform=self.train_transform)
+        self.test_set = datasets.__dict__[dataset_name](root='./backend/image_data_uploads', train=False, download=True, transform=self.test_transform)
+
+
+    @classmethod
+    def fromDefault(cls, dataset_name: str, train_transform=None, test_transform=None, batch_size: int = 32, shuffle: bool = True):
+        return cls(dataset_name, train_transform, test_transform, batch_size, shuffle)
+    def createTrainDataset(self) -> DataLoader:
+        return DataLoader(self.train_set, batch_size = self.batch_size, shuffle = self.shuffle, drop_last=True)
+
+    def createTestDataset(self):
+        return DataLoader(self.test_set, batch_size = self.batch_size, shuffle = self.shuffle, drop_last=True)
+    def getCategoryList(self) -> list[str]:
+        return self.train_set.classes if hasattr(self.train_set, "classes") else []

training/training/routes/image/__init__.py

@@ -0,0 +1,5 @@
+from training.routes.image.image import router
+
+
+def get_image_router():
+    return router

training/training/routes/image/image.py

@@ -0,0 +1,53 @@
+from typing import Literal, Optional
+from django.http import HttpRequest
+from ninja import Router, Schema
+from training.core.criterion import getCriterionHandler
+from training.core.dl_model import DLModel
+from training.core.datasets import ImageDefaultDatasetCreator
+from torch.utils.data import DataLoader
+from training.core.optimizer import getOptimizer
+from training.core.trainer import ClassificationTrainer
+from training.routes.image.schemas import ImageParams
+from training.core.authenticator import FirebaseAuth
+
+router = Router()
+
+@router.post("", auth=FirebaseAuth())
+def imageTrain(request: HttpRequest, imageParams: ImageParams):
+    if imageParams.default:
+        dataCreator = ImageDefaultDatasetCreator.fromDefault(
+            imageParams.default, imageParams.test_size, imageParams.shuffle
+        )
+        train_loader = DataLoader(
+            dataCreator.createTrainDataset(),
+            batch_size=imageParams.batch_size,
+            shuffle=False,
+            drop_last=True,
+        )
+
+        test_loader = DataLoader(
+            dataCreator.createTestDataset(),
+            batch_size=imageParams.batch_size,
+            shuffle=False,
+            drop_last=True,
+        )
+
+        model = DLModel.fromLayerParamsList(imageParams.user_arch)
+        optimizer = getOptimizer(model, imageParams.optimizer_name, 0.05)
+        criterionHandler = getCriterionHandler(imageParams.criterion)
+        if imageParams.problem_type == "CLASSIFICATION":
+            trainer = ClassificationTrainer(
+                train_loader,
+                test_loader,
+                model,
+                optimizer,
+                criterionHandler,
+                imageParams.epochs,
+                dataCreator.getCategoryList(),
+            )
+            for epoch_result in trainer:
+                print(epoch_result)
+            print(trainer.labels_last_epoch, trainer.y_pred_last_epoch)
+            print(trainer.generate_confusion_matrix())
+            print(trainer.generate_AUC_ROC_CURVE())
+            return trainer.generate_AUC_ROC_CURVE()

training/training/routes/image/schemas.py

@@ -0,0 +1,22 @@
+from typing import Any, Literal, Optional
+from ninja import Schema
+
+
+class LayerParams(Schema):
+    value: str
+    parameters: list[Any]
+
+
+class ImageParams(Schema):
+    target: str
+    features: list[str]
+    name: str
+    problem_type: Literal["CLASSIFICATION"]
+    default: Optional[str]
+    criterion: str 
+    optimizer_name: str
+    shuffle: bool
+    epochs: int
+    test_size: float
+    batch_size: int
+    user_arch: list[LayerParams]

training/training/urls.py

@@ -21,6 +21,7 @@
 from ninja import NinjaAPI, Schema
 from training.routes.datasets.default import get_default_datasets_router
 from training.routes.tabular import get_tabular_router
+from training.routes.image import get_image_router
 
 api = NinjaAPI()
 
@@ -32,6 +33,7 @@ def test(request: HttpRequest):
 
 api.add_router("/datasets/default/", get_default_datasets_router())
 api.add_router("/tabular", get_tabular_router())
+api.add_router("/image", get_image_router())
 
 urlpatterns = [
     path("admin/", admin.site.urls),