Setup package

.gitmodules

@@ -1,3 +1,3 @@
 [submodule "DeepLearningTemplate"]
 	path = DeepLearningTemplate
-	url = https://github.com/fastyangmh/DeepLearningTemplate
+	url = https://github.com/Heegreis/DeepLearningTemplate

README.md

@@ -5,6 +5,7 @@ This repository is used for image classification task and provides training, pre
 
 ```bash
 pip install -r requirements.txt
+pip install -e .
 ```
 
 ## Prepare dataset

src/data_preparation.py → imageclassification/data_preparation.py

@@ -1,6 +1,6 @@
 # import
-from src.project_parameters import ProjectParameters
-from DeepLearningTemplate.data_preparation import MyMNIST, MyCIFAR10, MyImageFolder, ImageLightningDataModule
+from imageclassification.project_parameters import ProjectParameters
+from deeplearningtemplate.data_preparation import MyMNIST, MyCIFAR10, MyImageFolder, ImageLightningDataModule
 
 
 # def

src/model.py → imageclassification/model.py

@@ -1,6 +1,6 @@
 # import
-from src.project_parameters import ProjectParameters
-from DeepLearningTemplate.model import BaseModel, load_from_checkpoint_for_supervised_model
+from imageclassification.project_parameters import ProjectParameters
+from deeplearningtemplate.model import BaseModel, load_from_checkpoint_for_supervised_model
 import timm
 from os.path import isfile
 import torch.nn as nn

src/predict.py → imageclassification/predict.py

@@ -1,9 +1,9 @@
 # import
-from src.project_parameters import ProjectParameters
-from DeepLearningTemplate.predict import ImagePredictDataset
-from src.model import create_model
+from imageclassification.project_parameters import ProjectParameters
+from deeplearningtemplate.predict import ImagePredictDataset
+from imageclassification.model import create_model
 import torch
-from DeepLearningTemplate.data_preparation import parse_transforms
+from deeplearningtemplate.data_preparation import parse_transforms
 from typing import Any
 from os.path import isfile
 from PIL import Image
@@ -29,7 +29,14 @@ def __init__(self, project_parameters) -> None:
 
     def predict(self, inputs) -> Any:
         result = []
-        if isfile(path=inputs):
+        if type(inputs) == Image.Image:
+            # the transformed sample dimension is (1, in_chans, width, height)
+            sample = self.transform(inputs)[None]
+            if self.device == 'cuda' and torch.cuda.is_available():
+                sample = sample.cuda()
+            with torch.no_grad():
+                result.append(self.model(sample).tolist()[0])
+        elif isfile(path=inputs):
             # predict the file
             sample = self.loader(inputs).convert(self.color_space)
             # the transformed sample dimension is (1, in_chans, width, height)
@@ -57,8 +64,6 @@ def predict(self, inputs) -> Any:
                         sample = sample.cuda()
                     result.append(self.model(sample).tolist())
         result = np.concatenate(result, 0)
-        print(', '.join(self.classes))
-        print(result)
         return result
 
 

src/predict_gui.py → imageclassification/predict_gui.py

@@ -1,9 +1,9 @@
 # import
-from src.project_parameters import ProjectParameters
-from DeepLearningTemplate.predict_gui import BasePredictGUI
-from src.predict import Predict
+from imageclassification.project_parameters import ProjectParameters
+from deeplearningtemplate.predict_gui import BasePredictGUI
+from imageclassification.predict import Predict
 from PIL import Image, ImageTk
-from DeepLearningTemplate.data_preparation import parse_transforms
+from deeplearningtemplate.data_preparation import parse_transforms
 from tkinter import Label, messagebox
 import tkinter as tk
 import gradio as gr

src/project_parameters.py → imageclassification/project_parameters.py

@@ -1,8 +1,7 @@
 # import
 import os
 import sys
-sys.path.append(os.path.realpath('DeepLearningTemplate/'))
-from DeepLearningTemplate.project_parameters import ProjectParameters
+from deeplearningtemplate.project_parameters import ProjectParameters
 
 if __name__ == '__main__':
     # project parameters

src/train.py → imageclassification/train.py

@@ -1,8 +1,8 @@
 # import
-from src.project_parameters import ProjectParameters
-from DeepLearningTemplate.train import BaseTrain
-from src.data_preparation import create_datamodule
-from src.model import create_model
+from imageclassification.project_parameters import ProjectParameters
+from deeplearningtemplate.train import BaseTrain
+from imageclassification.data_preparation import create_datamodule
+from imageclassification.model import create_model
 from typing import Any
 from shutil import copy2
 

src/tuning.py → imageclassification/tuning.py

@@ -1,6 +1,6 @@
 # import
-from src.project_parameters import ProjectParameters
-from DeepLearningTemplate.tuning import BaseTuning
+from imageclassification.project_parameters import ProjectParameters
+from deeplearningtemplate.tuning import BaseTuning
 from ray import tune
 from copy import deepcopy
 from typing import Any
@@ -10,7 +10,7 @@
 import ray
 from functools import partial
 from os.path import join
-from src.train import Train
+from imageclassification.train import Train
 
 
 # class

main.py

@@ -1,8 +1,8 @@
 # import
-from src.project_parameters import ProjectParameters
-from src.train import Train
-from src.predict import Predict
-from src.tuning import Tuning
+from imageclassification.project_parameters import ProjectParameters
+from imageclassification.train import Train
+from imageclassification.predict import Predict
+from imageclassification.tuning import Tuning
 
 
 # def
@@ -21,7 +21,7 @@ def main():
         result = Predict(project_parameters=project_parameters).predict(
             inputs=project_parameters.root)
     elif project_parameters.mode == 'predict_gui':
-        from src.predict_gui import PredictGUI
+        from imageclassification.predict_gui import PredictGUI
         result = PredictGUI(project_parameters=project_parameters).run()
     elif project_parameters.mode == 'tuning':
         result = Tuning(project_parameters=project_parameters,

requirements.txt

@@ -7,6 +7,7 @@ pytorch_lightning==1.4.9
 ray[tune]==1.9.1
 ruamel.yaml==0.17.17
 scikit_learn==1.0.2
+scipy==1.8.0
 seaborn==0.11.2
 timm==0.4.12
 torch==1.10.1
@@ -15,3 +16,4 @@ torchmetrics==0.7.1
 torchsummary==1.5.1
 torchvision==0.11.2
 tqdm==4.62.3
+albumentations==1.3.0

setup.py

@@ -0,0 +1,7 @@
+from setuptools import setup, find_packages
+
+setup(
+    name='ImageClassification',
+    version='0.0.1',
+    packages=find_packages(),
+)