├── DataLoader ├── classification_data_Loader.py ├── custom_data_Loader.py └── generated_data_Loader.py ├── README.md ├── cnn_inference.py ├── cnn_train.py ├── config.py ├── log_data.py ├── main.py ├── model_training ├── early_stop.py ├── model_architecture.py └── model_trainer.py ├── requirements.txt ├── test_accuracy.py └── user_inputs.ini /DataLoader/classification_data_Loader.py: -------------------------------------------------------------------------------- 1 | import os 2 | import numpy as np 3 | from tqdm import tqdm 4 | import cv2 5 | import torch 6 | from torch.utils.data import Dataset,DataLoader 7 | from torchvision import transforms 8 | 9 | 10 | def generate_data(DATA_FOLDER,IMAGE_SIZE=224,TEST_SIZE=0.1,SAVE_DATA=True,CLASSES_FILE="classes.txt",OUT_DATA="../generated_data"): 11 | 12 | if not os.path.exists(OUT_DATA): 13 | os.makedirs(OUT_DATA) 14 | 15 | classes=os.listdir(DATA_FOLDER) 16 | full_path=list(os.path.join(DATA_FOLDER,i) for i in classes) 17 | with open(os.path.join(OUT_DATA,CLASSES_FILE), 'w') as f: 18 | for i,data in enumerate(classes): 19 | f.write("%s\n" % data) 20 | 21 | data_list=[] 22 | labels = {full_path[x]:x for x in range(len(full_path))} 23 | for label in labels: 24 | for f in tqdm(os.listdir(label)): 25 | try: 26 | path=os.path.join(label,f) 27 | img=cv2.imread(path) 28 | img=cv2.resize(img,(IMAGE_SIZE,IMAGE_SIZE)) 29 | data_list.append([np.array(img),labels[label]]) 30 | 31 | except Exception as e: 32 | pass 33 | 34 | np.random.shuffle(data_list) 35 | test_size=int(len(data_list)*TEST_SIZE) 36 | training_data,testing_data=data_list[test_size:],data_list[:test_size] 37 | if SAVE_DATA: 38 | np.save(os.path.join(OUT_DATA,"training_data2.npy"),training_data) 39 | np.save(os.path.join(OUT_DATA,"testing_data2.npy"),testing_data) 40 | 41 | 42 | return training_data,testing_data 43 | 44 | 45 | class Classification_DATASET(Dataset): 46 | def __init__(self,data_set,transform=None): 47 | self.data_set=data_set 48 | self.transform=transform 49 | 50 | def __len__(self): 51 | return(len(self.data_set)) 52 | 53 | def __getitem__(self,idx): 54 | data = self.data_set[idx][0] 55 | data=data.astype(np.float32).reshape(3,100,100) 56 | 57 | if self.transform: 58 | data=self.transform(data) 59 | return (data,self.data_set[idx][1]) 60 | else: 61 | return (data,self.data_set[idx][1]) 62 | 63 | data_transform = transforms.Compose([ 64 | transforms.ToTensor(), 65 | transforms.Normalize(mean=[0.485, 0.456, 0.406], 66 | std=[0.229, 0.224, 0.225]) 67 | ]) 68 | 69 | 70 | training_data,testing_data=generate_data("../Data") 71 | 72 | 73 | train_loader = DataLoader(Classification_DATASET(training_data), batch_size=2, shuffle=False) 74 | test_loader = DataLoader(Classification_DATASET(testing_data), batch_size=2, shuffle=True) 75 | -------------------------------------------------------------------------------- /DataLoader/custom_data_Loader.py: -------------------------------------------------------------------------------- 1 | import os 2 | import numpy as np 3 | from tqdm import tqdm 4 | import cv2 5 | import argparse 6 | import torch 7 | from torch.utils.data import Dataset,DataLoader 8 | from torchvision import transforms 9 | import configparser 10 | 11 | config = configparser.ConfigParser() 12 | config.read('../user_inputs.ini') 13 | 14 | parser=argparse.ArgumentParser() 15 | parser.add_argument("--DATA_FOLDER", type=str, default=config.get('DataLoader','DATA_FOLDER'), help="Training Images folder") 16 | parser.add_argument("--IMAGE_SIZE", type=str, default=config.getint('DataLoader','IMAGE_SIZE'), help="Image Size") 17 | parser.add_argument("--TEST_SIZE", type=float, default=config.getfloat('DataLoader','TEST_SIZE'), help="Testing data percentage") 18 | parser.add_argument("--SAVE_DATA", type=bool, default=config.getboolean('DataLoader','SAVE_DATA'), help="Save data to numpy array") 19 | parser.add_argument("--SAVE_DATA_FOLDER", type=str, default=config.get('DataLoader','SAVE_DATA_FOLDER'), help="Save data to numpy array") 20 | parser.add_argument("--CLASSES_FILE", type=str, default="classes.txt", help="output text file contain all classes") 21 | parser.add_argument("--TESTING_DATA", type=str, default="training_data.npy", help="output numpy training data file name") 22 | parser.add_argument("--TRAINING_DATA", type=str, default="testing_data.npy", help="output numpy testing data file name") 23 | args = parser.parse_args() 24 | 25 | if not os.path.exists(args.SAVE_DATA_FOLDER): 26 | os.makedirs(args.SAVE_DATA_FOLDER) 27 | 28 | def generate_data(DATA_FOLDER,IMAGE_SIZE,TEST_SIZE): 29 | 30 | classes=os.listdir(DATA_FOLDER) 31 | full_path=list(os.path.join(DATA_FOLDER,i) for i in classes) 32 | with open(os.path.join(args.SAVE_DATA_FOLDER,args.CLASSES_FILE), 'w') as f: 33 | for i,data in enumerate(classes): 34 | f.write("%s\n" % data) 35 | 36 | data_list=[] 37 | labels = {full_path[x]:x for x in range(len(full_path))} 38 | for label in labels: 39 | for f in tqdm(os.listdir(label)): 40 | try: 41 | path=os.path.join(label,f) 42 | img=cv2.imread(path) 43 | img=cv2.resize(img,(IMAGE_SIZE,IMAGE_SIZE)) 44 | data_list.append([np.array(img),labels[label]]) 45 | 46 | except Exception as e: 47 | pass 48 | 49 | np.random.shuffle(data_list) 50 | test_size=int(len(data_list)*TEST_SIZE) 51 | training_data,testing_data=data_list[test_size:],data_list[:test_size] 52 | 53 | return training_data,testing_data 54 | 55 | 56 | training_data,testing_data=generate_data(args.DATA_FOLDER,args.IMAGE_SIZE,args.TEST_SIZE) 57 | 58 | if args.SAVE_DATA: #store data in a numpy array 59 | np.save(os.path.join(args.SAVE_DATA_FOLDER,args.TRAINING_DATA),training_data) #Save training data to numpy array 60 | np.save(os.path.join(args.SAVE_DATA_FOLDER,args.TESTING_DATA),testing_data) #Save testing data to numpy array 61 | 62 | 63 | class Classification_DATASET(Dataset): 64 | def __init__(self,data_set,transform=None): 65 | self.data_set=data_set 66 | self.transform=transform 67 | #self.labels=labels 68 | 69 | def __len__(self): 70 | return(len(self.data_set)) 71 | 72 | def __getitem__(self,idx): 73 | data = self.data_set[idx][0] 74 | data=data.astype(np.float32).reshape(3,100,100) 75 | 76 | if self.transform: 77 | data=self.transform(data) 78 | return (data,self.data_set[idx][1]) 79 | else: 80 | return (data,self.data_set[idx][1]) 81 | 82 | data_transform = transforms.Compose([ 83 | transforms.ToTensor(), 84 | transforms.Normalize(mean=[0.485, 0.456, 0.406], 85 | std=[0.229, 0.224, 0.225]) 86 | ]) 87 | 88 | 89 | train_data = Classification_DATASET(training_data) 90 | test_data = Classification_DATASET(testing_data) 91 | 92 | train_loader = DataLoader(train_data, batch_size=2, shuffle=False) 93 | test_loader = DataLoader(test_data, batch_size=2, shuffle=True) 94 | -------------------------------------------------------------------------------- /DataLoader/generated_data_Loader.py: -------------------------------------------------------------------------------- 1 | import os 2 | import numpy as np 3 | from tqdm import tqdm 4 | import cv2 5 | import argparse 6 | import torch 7 | from torch.utils.data import Dataset,DataLoader 8 | from torchvision import transforms 9 | import configparser 10 | 11 | config = configparser.ConfigParser() 12 | config.read('../user_inputs.ini') 13 | 14 | parser=argparse.ArgumentParser() 15 | parser.add_argument("--DATA_FOLDER", type=str, default=config.get('DataLoader','DATA_FOLDER'), help="Training Images folder") 16 | parser.add_argument("--IMAGE_SIZE", type=str, default=config.getint('DataLoader','IMAGE_SIZE'), help="Image Size") 17 | parser.add_argument("--TEST_SIZE", type=float, default=config.getfloat('DataLoader','TEST_SIZE'), help="Testing data percentage") 18 | parser.add_argument("--SAVE_DATA", type=bool, default=config.getboolean('DataLoader','SAVE_DATA'), help="Save data to numpy array") 19 | parser.add_argument("--SAVE_DATA_FOLDER", type=str, default=config.get('DataLoader','SAVE_DATA_FOLDER'), help="Save data to numpy array") 20 | parser.add_argument("--CLASSES_FILE", type=str, default="classes.txt", help="output text file contain all classes") 21 | parser.add_argument("--TESTING_DATA", type=str, default="training_data.npy", help="output numpy training data file name") 22 | parser.add_argument("--TRAINING_DATA", type=str, default="testing_data.npy", help="output numpy testing data file name") 23 | args = parser.parse_args() 24 | 25 | training_data=np.load(os.path.join(args.SAVE_DATA_FOLDER,args.TRAINING_DATA),allow_pickle=True) 26 | testing_data=np.load(os.path.join(args.SAVE_DATA_FOLDER,args.TESTING_DATA),allow_pickle=True) 27 | 28 | class Classification_DATASET(Dataset): 29 | def __init__(self,data_set,transform=None): 30 | self.data_set=data_set 31 | self.transform=transform 32 | #self.labels=labels 33 | 34 | def __len__(self): 35 | return(len(self.data_set)) 36 | 37 | def __getitem__(self,idx): 38 | data = self.data_set[idx][0] 39 | data=data.astype(np.float32).reshape(3,100,100) 40 | 41 | if self.transform: 42 | data=self.transform(data) 43 | return (data,self.data_set[idx][1]) 44 | else: 45 | return (data,self.data_set[idx][1]) 46 | 47 | data_transform = transforms.Compose([ 48 | transforms.ToTensor(), 49 | transforms.Normalize(mean=[0.485, 0.456, 0.406], 50 | std=[0.229, 0.224, 0.225]) 51 | ]) 52 | 53 | 54 | train_data = Classification_DATASET(training_data) 55 | test_data = Classification_DATASET(testing_data) 56 | 57 | print(train_data) 58 | 59 | train_loader = DataLoader(train_data, batch_size=2, shuffle=False) 60 | test_loader = DataLoader(test_data, batch_size=2, shuffle=True) 61 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # Simple CNN 2 | 3 | #### ```Simple CNN``` is a pipeline which can be use to train and infer CNN models by use of PyTorch and ONNX. It's simple and easy to USE !!! 🔥🔥 4 | ___ 5 | ### Install 6 | 7 | - Clone the repo and install **requirements.txt** in a Python environment 8 | 9 | ```bash 10 | git clone https://github.com/LahiRumesh/simple_cnn.git 11 | cd simple_cnn 12 | pip install -r requirements.txt 13 | ``` 14 | --- 15 | 16 | ### Data Preparation 17 | 18 | - Split images into **train** and **val** folders with each class the Image Folder 📂.. i.e for cat vs dogs classification, there should be a cat folder and dog folder in both train and val. The following folder structure illustrates 3 classes 19 | 20 | 21 | ```bash 22 | ├── Image_Folder 23 | ├── train 24 | │ │───── class1 25 | │ │ ├── class1.0.jpg 26 | │ │ ├── class1.1.jpg 27 | │ │ ├── class1.2.jpg 28 | │ │ ├── ......... 29 | │ │ └── class1.500.jpg 30 | │ │ 31 | │ │───── class2 32 | │ │ ├── class2.0.jpg 33 | │ │ ├── class2.1.jpg 34 | │ │ ├── class2.2.jpg 35 | │ │ ├── ......... 36 | │ │ └── class2.500.jpg 37 | │ │ 38 | │ └───── class3 39 | │ ├── class3.0.jpg 40 | │ ├── class3.1.jpg 41 | │ ├── class3.2.jpg 42 | │ ├── ......... 43 | │ └── class3.500.jpg 44 | │ 45 | └── val 46 | │───── class1 47 | │ ├── class1.501.jpg 48 | │ ├── class1.502.jpg 49 | │ ├── class1.503.jpg 50 | │ ├── ......... 51 | │ └── class1.600.jpg 52 | │ 53 | │───── class2 54 | │ ├── class2.501.jpg 55 | │ ├── class2.502.jpg 56 | │ ├── class2.503.jpg 57 | │ ├── ......... 58 | │ └── class2.600.jpg 59 | │ 60 | └───── class3 61 | ├── class3.501.jpg 62 | ├── class3.502.jpg 63 | ├── class3.503.jpg 64 | ├── ......... 65 | └── class3.600.jpg 66 | 67 | ``` 68 | 69 | --- 70 | 71 | ### Training 72 | #### After the data preparation, it's time for the training ! 73 | - Use the **config.py** to set the parameters, here are few parameters. 74 | 75 | ```bash 76 | cfg.data_dir = 'Data/Images/Image_Folder' # Image Folder path which contain train and val folders 77 | cfg.device = '0' # cuda device, i.e. 0 or 0,1,2,3 78 | cfg.image_size = 224 #input image size 79 | cfg.batch_size = 8 # batch size 80 | cfg.epochs = 50 #number of epochs 81 | 82 | cfg.model = 'resnet18' # torch vision classification model architectures for image classification 83 | # i.e. resnet18 or vgg16, alexnet, densenet121, squeezenet1_0 84 | 85 | cfg.pretrained = True # use pretrained weights for training 86 | 87 | #Early Stopping 88 | cfg.use_early_stopping = True # use Early stopping 89 | cfg.patience = 8 # how many epochs to wait before stopping when accuracy is not improving 90 | cfg.min_delta = 0.0001 # minimum difference between new accuracy and old accuracy for new accuracy to be considered as an improvement 91 | ``` 92 | 93 | 94 | - Here are the Available pre-trained models in ```Simple CNN``` 95 | 96 | 97 | 98 | 99 | 100 | 101 | 102 | 103 | 104 | 105 | 106 | 107 | 108 | 109 | 110 | 111 | 112 | 113 | 114 | 115 |
ArchitecturesAvailable Models
Resnetresnet18, resnet34, resnet50, resnet101, resnet152
VGGvgg13, vgg13_bn, vgg16, vgg16_bn, vgg19, vgg19_bn
Densenetdensenet121, densenet169, densenet161 , densenet201
Squeezenetsqueezenet1_0, squeezenet1_1
Alexnetalexnet
None: 9 | 10 | self.model_path = model_path 11 | self.img_size = img_size 12 | self.classes_path = classes_path 13 | model = onnx.load(model_path) 14 | self.session = rt.InferenceSession(model.SerializeToString()) 15 | 16 | def get_image(self, path, show=False): 17 | ''' 18 | Read the image and disply 19 | path : input image path 20 | show : display the image 21 | ''' 22 | img = cv2.imread(path) 23 | if show: 24 | cv2.imshow("Frame",img) 25 | cv2.waitKey(0) 26 | return img 27 | 28 | 29 | def preprocess(self, img, use_transform=False): 30 | 31 | ''' 32 | Image Pre-processing steps for inference 33 | img : image 34 | use_transform : use trasfromation step 35 | ''' 36 | img = img / 255. 37 | img = cv2.resize(img, (self.img_size, self.img_size)) 38 | if use_transform: 39 | img = (img - [0.485, 0.456, 0.406]) / [0.229, 0.224, 0.225] 40 | else: 41 | h, w = img.shape[0], img.shape[1] 42 | y0 = (h - 224) // 2 43 | x0 = (w - 224) // 2 44 | img = img[y0 : y0+224, x0 : x0+224, :] 45 | 46 | img = np.transpose(img, axes=[2,0, 1]) 47 | img = img.astype(np.float32) 48 | img = np.expand_dims(img, axis=0) 49 | return img 50 | 51 | def get_classes(self): 52 | 53 | ''' 54 | Read the class file and return class names as an array 55 | ''' 56 | with open(self.classes_path) as f: 57 | class_names = f.readlines() 58 | class_names = [c.strip() for c in class_names] 59 | return class_names 60 | 61 | 62 | def predict(self, path, 63 | show_img=False, 64 | use_transform=False): 65 | img = self.get_image(path, show=show_img) 66 | img = self.preprocess(img, use_transform=use_transform) 67 | inputs = {self.session.get_inputs()[0].name: img} 68 | preds = self.session.run(None, inputs)[0] 69 | preds = np.squeeze(preds) 70 | a = np.argsort(preds)[::-1] 71 | labels = self.get_classes() 72 | print('Predicted Class : %s' %(labels[a[0]])) 73 | 74 | 75 | 76 | 77 | if __name__ == '__main__': 78 | 79 | ''' 80 | python cnn_inference.py --model_path=models/cats_vs_dogs/cats_vs_dogs_resnet18_exp_1.onnx --class_path=models/cats_vs_dogs/classes.txt --img_path=test1.jpg --image_size=224 --use_transform=True 81 | 82 | ''' 83 | 84 | parser = argparse.ArgumentParser() 85 | parser.add_argument('--model_path', type=str, default='models/cats_vs_dogs/cats_vs_dogs_resnet18_exp_1.onnx', help='ONNX model path') 86 | parser.add_argument('--class_path', type=str, default='models/cats_vs_dogs/classes.txt', help='Class file path which contain class names') 87 | parser.add_argument('--img_path', type=str, default='test1.jpg', help='Input Image path') 88 | parser.add_argument('--image_size', type=int, default=224, help='Input Image size (Used for the training)') 89 | parser.add_argument('--show_image', type=bool, default=True, help='Display the image') 90 | parser.add_argument('--use_transform', type=bool, default=True, help='Use image transforms in pre-processing step') 91 | args = parser.parse_args() 92 | 93 | cnn_infer = CNNInference(args.model_path,args.image_size,args.class_path) 94 | cnn_infer.predict(args.img_path,show_img=args.show_image,use_transform=args.use_transform) 95 | 96 | -------------------------------------------------------------------------------- /cnn_train.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | import torch.optim as optim 4 | from torch.optim import lr_scheduler 5 | import torch.backends.cudnn as cudnn 6 | import torchvision 7 | from torchvision import datasets, models, transforms 8 | import os 9 | import copy 10 | import torch.onnx as torch_onnx 11 | from torch.autograd import Variable 12 | from config import cfg 13 | import wandb 14 | from tqdm import tqdm 15 | from time import sleep 16 | from log_data import model_configs, wandb_logs, get_accuracy 17 | from model_training.early_stop import EarlyStopping 18 | from model_training.model_architecture import ModelArchitecture 19 | cudnn.benchmark = True 20 | wandb.login() #login to wandb account 21 | 22 | 23 | class CNN_Trainer(): 24 | 25 | def __init__(self,image_dir, 26 | data_transforms) -> None: 27 | 28 | ''' 29 | load the images from the folder and use data transforms 30 | 31 | image_dir : path of the image folder 32 | data_transforms : data transforms 33 | 34 | ''' 35 | self.image_dir = image_dir 36 | self.data_transforms = data_transforms 37 | self.image_datasets = {x: datasets.ImageFolder(os.path.join(self.image_dir, x), 38 | self.data_transforms[x]) 39 | for x in ['train', 'val']} 40 | 41 | self.dataset_sizes = {x: len(self.image_datasets[x]) for x in ['train', 'val']} 42 | self.class_names = self.image_datasets['train'].classes 43 | self.checkpoints_dir = os.path.join("models", os.path.basename(self.image_dir)) 44 | 45 | try: 46 | os.makedirs(self.checkpoints_dir, exist_ok=True) 47 | except OSError: 48 | pass 49 | 50 | with open(os.path.join(self.checkpoints_dir,'classes.txt'), 'w') as f: 51 | for i,data in enumerate(self.class_names): 52 | f.write("%s\n" % data) 53 | 54 | def train_model(self,model, criterion, 55 | optimizer, scheduler, num_epochs=25, 56 | batch_size=4, 57 | shuffle=True,num_workers=4,patience=10,min_delta=0,use_early_stopping=True): 58 | 59 | 60 | dataloaders = {x: torch.utils.data.DataLoader(self.image_datasets[x], batch_size=batch_size, 61 | shuffle=shuffle, num_workers=num_workers) 62 | for x in ['train', 'val']} 63 | 64 | device = torch.device(f"cuda:{cfg.device}" if torch.cuda.is_available() else "cpu") 65 | inputs, classes = next(iter(dataloaders['train'])) 66 | out = torchvision.utils.make_grid(inputs) 67 | 68 | # Loard pre-trained model 69 | model_ = ModelArchitecture(model=cfg.model) 70 | model = model_.getModel(model,len(self.class_names)) 71 | model = model.to(device) 72 | 73 | # wandb initialization 74 | wandb.init(project=f"simple_cnn-{os.path.basename(self.image_dir)}".replace("/", "-"), config=cfg) 75 | images = wandb.Image(out, caption=f"Sample_Batch-{os.path.basename(self.image_dir)}") 76 | wandb.watch(model, criterion, log="all", log_freq=10) 77 | wandb.log({"images": images 78 | }) 79 | 80 | model_configs(cfg=cfg,classes=self.class_names) 81 | early_stopping = EarlyStopping(patience=cfg.patience, min_delta=cfg.min_delta) 82 | best_model_wts = copy.deepcopy(model.state_dict()) 83 | best_acc = 0.0 84 | 85 | for epoch in range(num_epochs): 86 | print(f'Epoch {epoch+1}') 87 | 88 | for phase in ['train', 'val']: 89 | if phase == 'train': 90 | model.train() # training mode 91 | else: 92 | model.eval() # evaluate mode 93 | 94 | running_loss = 0.0 95 | running_corrects = 0 96 | 97 | # Iterate over data. 98 | for inputs, labels in tqdm(dataloaders[phase]): 99 | sleep(0.01) 100 | inputs = inputs.to(device) 101 | labels = labels.to(device) 102 | optimizer.zero_grad() 103 | with torch.set_grad_enabled(phase == 'train'): 104 | outputs = model(inputs) 105 | _, preds = torch.max(outputs, 1) 106 | loss = criterion(outputs, labels) 107 | 108 | if phase == 'train': 109 | loss.backward() 110 | optimizer.step() 111 | 112 | # statistics and wandb logs 113 | running_loss += loss.item() * inputs.size(0) 114 | running_corrects += torch.sum(preds == labels.data) 115 | if phase == 'train': 116 | scheduler.step() 117 | 118 | epoch_loss = running_loss / self.dataset_sizes[phase] 119 | epoch_acc = running_corrects.double() / self.dataset_sizes[phase] 120 | 121 | get_accuracy(phase,epoch_loss,epoch_acc) 122 | wandb_logs(phase,epoch_loss,epoch_acc) # wandb log loss and accuracy 123 | 124 | # deep copy the model 125 | 126 | if phase == 'val': 127 | 128 | if cfg.use_early_stopping: 129 | early_stopping(epoch_acc,epoch) 130 | 131 | if epoch_acc > best_acc: 132 | best_acc = epoch_acc 133 | best_model_wts = copy.deepcopy(model.state_dict()) 134 | 135 | if cfg.use_early_stopping and early_stopping.early_stop: 136 | break 137 | print() 138 | 139 | 140 | # load best model weights and return for export 141 | model.load_state_dict(best_model_wts) 142 | return model 143 | 144 | 145 | 146 | def onnx_export(self,model,img_size=224,c_in=3,): 147 | 148 | ''' 149 | Export the onnx model from the best weight file, 150 | 151 | model : trained model 152 | img_size : train image size 153 | c_in : colour channels (r,g,b) 154 | 155 | ''' 156 | 157 | input_shape = (c_in, img_size, img_size) 158 | model_prefix = os.path.basename(self.image_dir) + "_" + cfg.model 159 | counter = 1 160 | 161 | while os.path.exists(os.path.join(self.checkpoints_dir, f'{model_prefix + "_exp_" + str(counter)}.onnx')): 162 | counter += 1 163 | 164 | # 165 | # if you are using cpu for training use dummy input as 166 | # dummy_input = Variable(torch.randn(1, *input_shape)) 167 | # 168 | 169 | dummy_input = Variable(torch.randn(1, *input_shape,device="cuda")) 170 | torch_onnx.export(model, 171 | dummy_input, 172 | os.path.join(self.checkpoints_dir,f'{model_prefix + "_exp_" + str(counter)}.onnx'), 173 | verbose=False) 174 | print(f"{model_prefix}_exp_" + str(counter)+ ".onnx model successfully exported !") 175 | 176 | 177 | if __name__ == '__main__': 178 | 179 | 180 | data_transforms = { 181 | 'train': transforms.Compose([ 182 | transforms.Resize((cfg.image_size,cfg.image_size)), 183 | transforms.ToTensor(), 184 | transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) 185 | ]), 186 | 'val': transforms.Compose([ 187 | transforms.Resize((cfg.image_size,cfg.image_size)), 188 | transforms.ToTensor(), 189 | transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) 190 | ]), 191 | } 192 | 193 | train_CNN = CNN_Trainer(cfg.data_dir,data_transforms) 194 | 195 | # model initilization 196 | model = getattr(models, cfg.model)(pretrained=cfg.pretrained) 197 | loss_criterion = getattr(nn, cfg.loss_criterion)() 198 | optimizer = getattr(optim, cfg.optimizer)(model.parameters(), 199 | lr=cfg.learning_rate, momentum=cfg.momentum, 200 | weight_decay=cfg.weight_decay) 201 | 202 | # Decay Learning Rate 203 | exp_lr_scheduler = getattr(lr_scheduler, cfg.lr_scheduler)(optimizer, step_size=cfg.steps, gamma=cfg.gamma) 204 | 205 | cnn_model = train_CNN.train_model(model,loss_criterion,optimizer,exp_lr_scheduler,num_epochs=cfg.epochs,batch_size=cfg.batch_size) 206 | 207 | # export the ONNX model 208 | 209 | train_CNN.onnx_export(cnn_model,img_size=cfg.image_size) 210 | -------------------------------------------------------------------------------- /config.py: -------------------------------------------------------------------------------- 1 | from easydict import EasyDict 2 | 3 | 4 | cfg = EasyDict() 5 | 6 | cfg.data_dir = 'data_dir' # Image Folder path which contain train and val folders 7 | cfg.device = '0' # cuda device, i.e. 0 or 0,1,2,3 8 | 9 | cfg.image_size = 224 #input image size 10 | cfg.batch_size = 4 # batch size 11 | cfg.epochs = 40 #number of epochs 12 | cfg.learning_rate = 0.001 13 | cfg.momentum = 0.9 14 | cfg.weight_decay = 0 15 | 16 | 17 | #model configuration 18 | cfg.model = 'resnet50' # model architecture 19 | cfg.pretrained = True # Use pretrained weight 20 | cfg.loss_criterion = 'CrossEntropyLoss' # Loss Function 21 | cfg.optimizer = 'SGD' # optimizer 22 | cfg.lr_scheduler = 'StepLR' # Learning Rate Decay 23 | cfg.steps = 7 # Decay epochs 24 | cfg.gamma = 0.1 # Decay Learning Rate factor 25 | 26 | #Early Stopping 27 | cfg.use_early_stopping = True #Use Early stopping 28 | cfg.patience = 8 # how many epochs to wait before stopping when accuracy is not improving 29 | cfg.min_delta = 0.0001 # minimum difference between new accuracy and old accuracy for new accuracy to be considered as an improvement -------------------------------------------------------------------------------- /log_data.py: -------------------------------------------------------------------------------- 1 | import os 2 | from tabulate import tabulate 3 | import wandb 4 | 5 | 6 | 7 | def model_configs(cfg,classes): 8 | 9 | header = f"Training {os.path.basename(cfg.data_dir)} Data Set" 10 | print('-' * len(header)) 11 | print(header) 12 | print('-' * len(header)) 13 | config_table = [["Model",cfg.model], 14 | ["Batch Size",cfg.batch_size], 15 | ["Image Size",cfg.image_size], 16 | ["Epochs",cfg.epochs], 17 | ["Learning Rate",cfg.learning_rate], 18 | ["Use Pretrained",cfg.pretrained], 19 | ["Momentum",cfg.momentum], 20 | ["Weight Decay",cfg.weight_decay], 21 | ["Loss Function",cfg.loss_criterion], 22 | ["Optimizer",cfg.optimizer], 23 | ["Learning Rate Decay",cfg.lr_scheduler], 24 | ["Decay epochs ",cfg.steps], 25 | ["Decay Learning Rate factor ",cfg.optimizer], 26 | ["Classes", str(classes)], 27 | ["Number of Classes", len(classes)]] 28 | print(tabulate(config_table,tablefmt="fancy_grid")) 29 | 30 | 31 | 32 | def wandb_logs(phase,epoch_loss,epoch_acc): 33 | 34 | if phase == 'train': 35 | wandb.log({"train_loss": epoch_loss, 36 | "train_accuracy" : epoch_acc 37 | }) 38 | 39 | elif phase == 'val': 40 | wandb.log({"val_loss": epoch_loss, 41 | "val_accuracy" : epoch_acc 42 | }) 43 | 44 | 45 | def get_accuracy(phase,epoch_loss,epoch_acc): 46 | header = f'{phase} Loss: {epoch_loss:.4f} Accuracy: {epoch_acc:.4f}' 47 | print('-' * len(header)) 48 | print(header) 49 | print('-' * len(header)) 50 | -------------------------------------------------------------------------------- /main.py: -------------------------------------------------------------------------------- 1 | from DataLoader.classification_data_Loader import generate_data,Classification_DATASET 2 | from torch.utils.data import DataLoader 3 | 4 | 5 | 6 | device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") 7 | 8 | training_data,testing_data=generate_data("Data") 9 | 10 | train_loader = DataLoader(Classification_DATASET(training_data), batch_size=2, shuffle=False) 11 | test_loader = DataLoader(Classification_DATASET(testing_data), batch_size=2, shuffle=True) 12 | 13 | #define NN 14 | -------------------------------------------------------------------------------- /model_training/early_stop.py: -------------------------------------------------------------------------------- 1 | 2 | class EarlyStopping(): 3 | """ 4 | Early stopping to stop the training when the loss accuracy not improve after 5 | certain epochs. 6 | """ 7 | def __init__(self, patience=5, min_delta=0): 8 | """ 9 | :param patience: how many epochs to wait before stopping when accuracy is 10 | not improving 11 | :param min_delta: minimum difference between new accuracy and old accuracy for 12 | new accuracy to be considered as an improvement 13 | """ 14 | self.patience = patience 15 | self.min_delta = min_delta 16 | self.counter = 0 17 | self.best_acc = None 18 | self.early_stop = False 19 | 20 | 21 | def __call__(self, val_acc,epoch): 22 | 23 | if self.best_acc == None: 24 | self.best_acc = val_acc 25 | elif val_acc - self.best_acc > self.min_delta: 26 | self.best_acc = val_acc 27 | elif val_acc - self.best_acc < self.min_delta: 28 | self.counter += 1 29 | print(f" Early stopping counter {self.counter} of {self.patience}") 30 | if self.counter >= self.patience: 31 | print('Early stopping...') 32 | self.early_stop = True -------------------------------------------------------------------------------- /model_training/model_architecture.py: -------------------------------------------------------------------------------- 1 | import torch.nn as nn 2 | from tabulate import tabulate 3 | 4 | class ModelArchitecture(): 5 | def __init__(self, model='resnet50') -> None: 6 | self.model = model 7 | self.resnet = ["resnet18","resnet34","resnet50","resnet101","resnet152"] 8 | self.vgg = ["vgg13","vgg13_bn","vgg16","vgg16_bn","vgg19","vgg19_bn"] 9 | self.alexnet = ["alexnet"] 10 | self.squeezenet = ["squeezenet1_0","squeezenet1_1"] 11 | self.densenet = ["densenet121", "densenet169", "densenet161" ,"densenet201"] 12 | 13 | 14 | def getModel(self, 15 | model, 16 | out_features): 17 | 18 | if self.model in self.resnet: 19 | num_ftrs = model.fc.in_features 20 | model.fc = nn.Linear(num_ftrs, out_features) 21 | 22 | elif self.model in self.vgg: 23 | num_ftrs = model.classifier[6].in_features 24 | model.classifier[6] = nn.Linear(num_ftrs,out_features) 25 | 26 | elif self.model in self.alexnet: 27 | num_ftrs = model.classifier[6].in_features 28 | model.classifier[6] = nn.Linear(num_ftrs,out_features) 29 | 30 | elif self.model in self.squeezenet: 31 | model.classifier[1] = nn.Conv2d(512, out_features, kernel_size=(1,1), stride=(1,1)) 32 | 33 | elif self.model in self.densenet: 34 | num_ftrs = model.classifier.in_features 35 | model.classifier = nn.Linear(num_ftrs, out_features) 36 | 37 | else: 38 | print('Available Models') 39 | model_table = [["Resnet",["resnet18","resnet34","resnet50","resnet101","resnet152"]], 40 | ["VGG",["vgg13","vgg13_bn","vgg16","vgg16_bn","vgg19","vgg19_bn"]], 41 | ["Densenet",["densenet121", "densenet169", "densenet161" ,"densenet201"]], 42 | ["Squeezenet",["squeezenet1_0","squeezenet1_1"]], 43 | ["Alexnet", ["alexnet"]]] 44 | headers = ["Architectures", "Available Models"] 45 | print(tabulate(model_table,headers,tablefmt="psql")) 46 | raise Exception("Invalid Model Architecture !! Please use available models") 47 | 48 | 49 | return model -------------------------------------------------------------------------------- /model_training/model_trainer.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch 3 | import torch.nn as nn 4 | import torch.optim as optim 5 | from torch.optim import lr_scheduler 6 | from torch.utils.data import Dataset,DataLoader 7 | from torchvision import transforms,models 8 | 9 | device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") 10 | 11 | 12 | class Simple_CNN_Trainer(): 13 | def __init__(self,model,learning_rate,num_epochs,train_loader,test_loader): 14 | self.model=model 15 | self.learning_rate=learning_rate 16 | self.num_epochs=num_epochs 17 | self.train_loader=train_loader 18 | self.test_loader=test_loader 19 | 20 | model_ft = models.resnet18(pretrained=True) 21 | num_ftrs = model_ft.fc.in_features 22 | model_ft.fc = nn.Linear(num_ftrs, len(classes)) 23 | 24 | model = model_ft.to(device) 25 | learning_rate = 0.1 26 | num_epochs = 5 27 | 28 | criterion = nn.CrossEntropyLoss() 29 | optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate) 30 | 31 | # Train the model 32 | 33 | total_step = len(train_loader) 34 | for epoch in range(num_epochs): 35 | for i, (images, labels) in enumerate(train_loader): 36 | images = images.to(device) 37 | labels = labels.to(device) 38 | print("Label Shape=",labels.shape) 39 | # Forward pass 40 | outputs = model(images) 41 | print("output Shape=",outputs.shape) 42 | loss = criterion(outputs, labels) 43 | 44 | # Backward and optimize 45 | optimizer.zero_grad() 46 | loss.backward() 47 | optimizer.step() 48 | 49 | if (i+1) % 2 == 0: 50 | print ('Epoch [{}/{}], Step [{}/{}], Loss: {:.4f}' 51 | .format(epoch+1, num_epochs, i+1, total_step, loss.item())) 52 | 53 | # Test the model 54 | model.eval() # eval mode (batchnorm uses moving mean/variance instead of mini-batch mean/variance) 55 | with torch.no_grad(): 56 | correct = 0 57 | total = 0 58 | for images, labels in test_loader: 59 | images = images.to(device) 60 | labels = labels.to(device) 61 | outputs = model(images) 62 | _, predicted = torch.max(outputs.data, 1) 63 | total += labels.size(0) 64 | correct += (predicted == labels).sum().item() 65 | 66 | # Save the model checkpoint 67 | torch.save(model.state_dict(), 'model.ckpt') -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- 1 | easydict==1.9 2 | imutils==0.5.4 3 | matplotlib==3.5.1 4 | onnx==1.11.0 5 | onnxruntime==1.11.0 6 | opencv-python==4.5.5.62 7 | packaging==21.3 8 | pandas==1.3.5 9 | Pillow==8.4.0 10 | psutil==5.9.0 11 | pytz==2021.3 12 | PyYAML==6.0 13 | pyzmq==19.0.2 14 | regex==2021.11.10 15 | requests==2.26.0 16 | requests-oauthlib==1.3.0 17 | setproctitle==1.2.2 18 | tabulate==0.8.9 19 | timm==0.4.12 20 | torch==1.10.1 21 | torchmetrics==0.6.2 22 | torchvision==0.11.2 23 | tqdm==4.62.3 24 | urllib3==1.26.7 25 | wandb==0.12.11 26 | zipp==3.7.0 27 | -------------------------------------------------------------------------------- /test_accuracy.py: -------------------------------------------------------------------------------- 1 | import os 2 | import onnx 3 | import numpy as np 4 | import onnxruntime as rt 5 | import cv2 6 | import pandas as pd 7 | import argparse 8 | from tqdm import tqdm 9 | import matplotlib.pyplot as plt 10 | 11 | class CalAccuracy(): 12 | 13 | def __init__(self,model_path,img_size,classes_path) -> None: 14 | 15 | self.model_path = model_path 16 | self.img_size = img_size 17 | self.classes_path = classes_path 18 | model = onnx.load(model_path) 19 | self.session = rt.InferenceSession(model.SerializeToString()) 20 | 21 | 22 | def getImageList(self, 23 | dirName, 24 | endings=['.jpg','.jpeg','.png','.JPG']): 25 | 26 | listOfFile = os.listdir(dirName) 27 | allFiles = list() 28 | 29 | for i,ending in enumerate(endings): 30 | if ending[0]!='.': 31 | endings[i] = '.'+ending 32 | # Iterate over all the entries 33 | for entry in listOfFile: 34 | # Create full path 35 | fullPath = os.path.join(dirName, entry) 36 | # If entry is a directory then get the list of files in this directory 37 | if os.path.isdir(fullPath): 38 | allFiles = allFiles + self.getImageList(fullPath,endings) 39 | else: 40 | for ending in endings: 41 | if entry.endswith(ending): 42 | allFiles.append(fullPath) 43 | return allFiles#, os.path.basename(dirName) 44 | 45 | 46 | 47 | def preprocess(self, img, use_transform=False): 48 | 49 | ''' 50 | Image Pre-processing steps for inference 51 | img : image 52 | use_transform : use trasfromation step 53 | ''' 54 | img = img / 255. 55 | img = cv2.resize(img, (self.img_size, self.img_size)) 56 | if use_transform: 57 | img = (img - [0.485, 0.456, 0.406]) / [0.229, 0.224, 0.225] 58 | else: 59 | h, w = img.shape[0], img.shape[1] 60 | y0 = (h - 224) // 2 61 | x0 = (w - 224) // 2 62 | img = img[y0 : y0+224, x0 : x0+224, :] 63 | 64 | img = np.transpose(img, axes=[2,0, 1]) 65 | img = img.astype(np.float32) 66 | img = np.expand_dims(img, axis=0) 67 | return img 68 | 69 | def get_classes(self): 70 | 71 | ''' 72 | Read the class file and return class names as an array 73 | ''' 74 | with open(self.classes_path) as f: 75 | class_names = f.readlines() 76 | class_names = [c.strip() for c in class_names] 77 | return class_names 78 | 79 | 80 | def predict(self, dir_path, 81 | use_transform=False, 82 | test_dir='test_results'): 83 | 84 | ''' 85 | Infer from the Test data set and visualize results 86 | ''' 87 | 88 | counter = 1 89 | while os.path.exists(os.path.join(test_dir, f'{"test_exp_" + str(counter)}')): 90 | counter += 1 91 | 92 | try: 93 | os.makedirs(os.path.join(test_dir, f'{"test_exp_" + str(counter)}'), exist_ok=True) 94 | except OSError: 95 | pass 96 | 97 | df = pd.DataFrame() 98 | img_name, true_class, predict_class = [], [], [] 99 | input_paths = self.getImageList(dir_path) 100 | for img_ in tqdm(input_paths): 101 | dir_class = os.path.basename(os.path.dirname(img_)) 102 | img = cv2.imread(img_) 103 | img = self.preprocess(img, use_transform=use_transform) 104 | inputs = {self.session.get_inputs()[0].name: img} 105 | preds = self.session.run(None, inputs)[0] 106 | preds = np.squeeze(preds) 107 | a = np.argsort(preds)[::-1] 108 | labels = self.get_classes() 109 | img_name.append(img_) 110 | true_class.append(dir_class) 111 | predict_class.append(labels[a[0]]) 112 | 113 | df["Image"] = img_name 114 | df["True_Class"] = true_class 115 | df["Predicted_Class"] = predict_class 116 | df.to_csv(os.path.join(os.path.join(test_dir, f'{"test_exp_" + str(counter)}'),"predict_results.csv")) 117 | 118 | df['result'] = df['True_Class'] == df['Predicted_Class'] 119 | mean_accuracy = df.value_counts('result')[True] / df['result'].count() 120 | df2 = pd.concat([df.groupby('True_Class')['result'].sum().reset_index(),df.value_counts('True_Class').reset_index()],axis=1,join='inner') 121 | df2['Accuracy'] = df2.iloc[:,1]/df2.iloc[:,3] 122 | df2 = df2.loc[:, ~df2.columns.duplicated()] 123 | df2 = df2.drop(columns=df2.columns.values[1:3]) 124 | df2.to_csv(os.path.join(os.path.join(test_dir, f'{"test_exp_" + str(counter)}'),"accuracy_results.csv")) 125 | print('Model Accuracy : %s' %(mean_accuracy)) 126 | df2.plot(x="True_Class", y='Accuracy', kind='bar') 127 | plt.savefig(os.path.join(os.path.join(test_dir, f'{"test_exp_" + str(counter)}'),"test_accuracy.png")) 128 | plt.show() 129 | 130 | 131 | 132 | 133 | 134 | if __name__ == '__main__': 135 | 136 | ''' 137 | python test_accuracy.py --model_path=models/cats_vs_dogs/cats_vs_dogs_resnet18_exp_1.onnx --class_path=models/cats_vs_dogs/classes.txt --img_dir=Classification/Data/cats_vs_dogs/test --image_size=224 --use_transform=True 138 | 139 | ''' 140 | 141 | parser = argparse.ArgumentParser() 142 | parser.add_argument('--model_path', type=str, default='models/cats_vs_dogs/cats_vs_dogs_resnet50_exp_8.onnx', help='ONNX model path') 143 | parser.add_argument('--class_path', type=str, default='models/cats_vs_dogs/classes.txt', help='Class file path which contain class names') 144 | parser.add_argument('--img_dir', type=str, default='Classification/Data/cats_vs_dogs/test', help='Test Images Dir path') 145 | parser.add_argument('--image_size', type=int, default=224, help='Input Image size (Used for the training)') 146 | parser.add_argument('--use_transform', type=bool, default=True, help='Use image transforms in pre-processing step') 147 | args = parser.parse_args() 148 | 149 | cal_accuracy = CalAccuracy(args.model_path,args.image_size,args.class_path) 150 | cal_accuracy.predict(args.img_dir,use_transform=args.use_transform) 151 | 152 | -------------------------------------------------------------------------------- /user_inputs.ini: -------------------------------------------------------------------------------- 1 | [DataLoader] 2 | DATA_FOLDER=../Data 3 | IMAGE_SIZE=100 4 | TEST_SIZE=0.1 5 | SAVE_DATA=True 6 | SAVE_DATA_FOLDER=../generated_data 7 | 8 | --------------------------------------------------------------------------------