Terragon Image Classifier (TIC-Torch)¶
based on a PyTorch Convolutional Neural Network¶
Why TIC-Torch?¶
The TIC-Torch application offers a convenient, fast and effective workflow to create, train and deploy individual and versatile image classifiers.
What's the workflow?¶
- Collect many jpg files of the objects you want to classify and arrange them in the "raw" directory. Create one subfolder for each object you want to classify and put the images inside. Example:
- raw/dogs
- raw/cats
- raw/humans
- Run the TIC-Torch app to resize and vary the images and to train the Neural Network.
- Run the TIC-Torch app to classify a new jpg.
Import OS and Image manipulation libraries¶
In [1]:
# System Librariesimportos,sysfromdistutils.dir_utilimportcopy_treefromPILimportImage,ImageEnhanceimportPIL.ImageOpsimportrandom# PyTorch Librariesimporttorchimporttorchvisionimporttorchvision.transformsastransformsimporttorch.nnasnnimporttorch.optimasoptimfromtorch.autogradimportVariableimporttorch.nn.functionalasF# Other Libraries we'll useimportnumpyasnpimportmatplotlib.pyplotasplt%matplotlib inline
print("Libraries imported - ready to use PyTorch",torch.__version__)
Define Variables¶
In [2]:
path_images_raw="raw/"path_images_resized="resized/"resize_x=128resize_y=128batch_size=50shuffle=Trueepochs=100
raw/¶
In [3]:
path_images_raw_subdirs=os.listdir(path_images_raw)print("raw subdirs: ",path_images_raw_subdirs)
resized/ copy¶
In [4]:
defImagesCopy(path_images_raw,path_images_resized):copy_tree(path_images_raw,path_images_resized)path_images_resized_subdirs=os.listdir(path_images_resized)print("copy from raw to resized folder")print("resized subdirs: ",path_images_resized_subdirs)ImagesCopy(path_images_raw,path_images_resized)
Resize¶
In [5]:
defImagesResize():forsubdir,dirs,filesinos.walk(path_images_resized):forfileinfiles:path_image=os.path.join(subdir,file)ifos.path.isfile(path_image):im=Image.open(path_image)f,e=os.path.splitext(path_image)imResize=im.resize((resize_x,resize_y),Image.ANTIALIAS)imResize.save(f+'_resized.jpg','JPEG',quality=90)os.remove(path_image)print("resized all images in folder resized/")ImagesResize()
Variations¶
In [6]:
defImageVariations():forsubdir,dirs,filesinos.walk(path_images_resized):forfileinfiles:path_image=os.path.join(subdir,file)ifos.path.isfile(path_image):im=Image.open(path_image)f,e=os.path.splitext(path_image)enhancer=ImageEnhance.Brightness(im)imLighter=enhancer.enhance(1.8)imLighter.save(f+'_lighter.jpg','JPEG',quality=90)imDarker=enhancer.enhance(0.5)imDarker.save(f+'_darker.jpg','JPEG',quality=90)imInverted=PIL.ImageOps.invert(im)imInverted.save(f+'_inverted.jpg','JPEG',quality=90)randomRotate=random.randrange(-45,45)imRotated=im.rotate(randomRotate)imRotated.save(f+'_rotated.jpg','JPEG',quality=90)imMirror=im.transpose(Image.FLIP_LEFT_RIGHT)imMirror.save(f+'_mirrored.jpg','JPEG',quality=90)imGrey=im.convert('L')imGrey.save(f+'_greyscale.jpg','JPEG',quality=90)print("added a mirror and greyscale images for all images in folder resized/")ImageVariations()
Images to Tensors¶
In [7]:
# Function to ingest data using training and test loaders# Now load the images from the resized folderdefload_dataset(path_images_resized):# Load all of the imagestransformation=transforms.Compose([# transform to tensorstransforms.ToTensor(),# Normalize the pixel values (in R, G, and B channels)transforms.Normalize(mean=[0.5,0.5,0.5],std=[0.5,0.5,0.5])])# Load all of the images, transforming themfull_dataset=torchvision.datasets.ImageFolder(root=path_images_resized,transform=transformation)# Split into training (70% and testing (30%) datasets)train_size=int(0.7*len(full_dataset))test_size=len(full_dataset)-train_sizetrain_dataset,test_dataset=torch.utils.data.random_split(full_dataset,[train_size,test_size])# define a loader for the training data we can iterate through in 50-image batchestrain_loader=torch.utils.data.DataLoader(train_dataset,batch_size=batch_size,num_workers=0,shuffle=shuffle)# define a loader for the testing data we can iterate through in 50-image batchestest_loader=torch.utils.data.DataLoader(test_dataset,batch_size=batch_size,num_workers=0,shuffle=shuffle)returntrain_loader,test_loader# Get the class namesclasses=os.listdir(path_images_resized)classes.sort()classes_amount=len(classes)print(len(classes),'classes:')print(classes)# Get the iterative dataloaders for test and training datatrain_loader,test_loader=load_dataset(path_images_resized)
Neural Net Model¶
In [8]:
# Create a neural net classclassNet(nn.Module):# Constructordef__init__(self,num_classes=classes_amount):super(Net,self).__init__()# Our images are RGB, so input channels = 3. wW'll apply 12 filters in the first convolutional layerself.conv1=nn.Conv2d(in_channels=3,out_channels=12,kernel_size=3,stride=1,padding=1)# We'll apply max pooling with a kernel size of 2self.pool=nn.MaxPool2d(kernel_size=2)# A second convolutional layer takes 12 input channels, and generates 12 outputsself.conv2=nn.Conv2d(in_channels=12,out_channels=12,kernel_size=3,stride=1,padding=1)# A third convolutional layer takes 12 inputs and generates 24 outputsself.conv3=nn.Conv2d(in_channels=12,out_channels=24,kernel_size=3,stride=1,padding=1)# A drop layer deletes 20% of the features to help prevent overfittingself.drop=nn.Dropout2d(p=0.2)# Our 128x128 image tensors will be pooled twice with a kernel size of 2. 128/2/2 is 32.# So our feature tensors are now 32 x 32, and we've generated 24 of them# We need to flatten these to map them to the probability for each classself.fc=nn.Linear(in_features=32*32*24,out_features=num_classes)defforward(self,x):# Use a relu activation function after layer 1 (convolution 1 and pool)x=F.relu(self.pool(self.conv1(x)))# Use a relu activation function after layer 1 (convolution 2 and drop)# Use a relu activation function after layer 3 (convolution 3)x=F.relu(self.pool(self.conv2(x)))# Drop some features after the 3rd convolution to prevent overfittingx=F.relu(self.drop(self.conv3(x)))# Only drop the features if this is a training passx=F.dropout(x,training=self.training)# Flattenx=x.view(-1,32*32*24)x=self.fc(x)# Return class probabilities via a softmax function returnF.log_softmax(x,dim=1)print("CNN model class defined!")
Train and Test¶
In [10]:
deftrain(model,device,train_loader,optimizer,epoch):# Set the model to training modemodel.train()train_loss=0print("Epoch:",epoch)# Process the images in batchesforbatch_idx,(data,target)inenumerate(train_loader):# Use the CPU or GPU as appropriatedata,target=data.to(device),target.to(device)# Reset the optimizeroptimizer.zero_grad()# Push the data forward through the model layersoutput=model(data)# Get the lossloss=loss_criteria(output,target)# Keep a running totaltrain_loss+=loss.item()# Backpropagateloss.backward()optimizer.step()# Print metrics for every 10 batches so we see some progressifbatch_idx%10==0:print('Training set [{}/{} ({:.0f}%)] Loss: {:.6f}'.format(batch_idx*len(data),len(train_loader.dataset),100.*batch_idx/len(train_loader),loss.item()))# return average loss for the epochreturntrain_loss/len(train_loader.dataset)deftest(model,device,test_loader):# Switch the model to evaluation mode (so we don't backpropagate or drop)model.eval()test_loss=0correct=0withtorch.no_grad():fordata,targetintest_loader:data,target=data.to(device),target.to(device)# Get the predicted classes for this batchoutput=model(data)# calculate the loss and successful predictions for this batchtest_loss+=loss_criteria(output,target).item()pred=output.max(1,keepdim=True)[1]correct+=pred.eq(target.view_as(pred)).sum().item()# Calculate the average loss and total accuracy for this epochtest_loss/=len(test_loader.dataset)print('Test set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\n'.format(test_loss,correct,len(test_loader.dataset),100.*correct/len(test_loader.dataset)))# return average loss for the epochreturntest_loss# Now use the train and test functions to train and test the model device="cpu"if(torch.cuda.is_available()):# if GPU available, use cuda (on a cpu, training will take a considerable length of time!)device="cuda"print('Training on',device)# Create an instance of the model class and allocate it to the devicemodel=Net(num_classes=len(classes)).to(device)# Use an "Adam" optimizer to adjust weights# (see https://pytorch.org/docs/stable/optim.html#algorithms for details of supported algorithms)optimizer=optim.Adam(model.parameters(),lr=0.001)# Specify the loss criterialoss_criteria=nn.CrossEntropyLoss()# Track metrics in these arraysepoch_nums=[]training_loss=[]validation_loss=[]# Train over 5 epochs (in a real scenario, you'd likely use many more)epochs=epochsifos.path.isfile("model.pt"):print("model.pt already exists and is now loaded, no training!")model=torch.load("model.pt")else:forepochinrange(1,epochs+1):train_loss=train(model,device,train_loader,optimizer,epoch)test_loss=test(model,device,test_loader)epoch_nums.append(epoch)training_loss.append(train_loss)validation_loss.append(test_loss)%matplotlib inline
frommatplotlibimportpyplotaspltplt.plot(epoch_nums,training_loss)plt.plot(epoch_nums,validation_loss)plt.xlabel('epoch')plt.ylabel('loss')plt.legend(['training','validation'],loc='upper right')plt.show()torch.save(model,"model.pt")
Evaluate - Confusion Matrix¶
In [11]:
#Pytorch doesn't have a built-in confusion matrix metric, so we'll use SciKit-Learnfromsklearn.metricsimportconfusion_matrix# Set the model to evaluate modemodel.eval()# Get predictions for the test data and convert to numpy arrays for use with SciKit-Learnprint("Getting predictions from test set...")truelabels=[]predictions=[]fordata,targetintest_loader:forlabelintarget.cpu().data.numpy():truelabels.append(label)forpredictioninmodel.cpu()(data).data.numpy().argmax(1):predictions.append(prediction)# Plot the confusion matrixcm=confusion_matrix(truelabels,predictions)plt.imshow(cm,interpolation="nearest",cmap=plt.cm.Greys)plt.colorbar()tick_marks=np.arange(len(classes))plt.xticks(tick_marks,classes,rotation=45)plt.yticks(tick_marks,classes)plt.xlabel("Predicted Shape")plt.ylabel("True Shape")plt.show()
Predict¶
In [13]:
# Function to predict the class of an imagedefpredict_image(classifier,image):importnumpy# Set the classifer model to evaluation modeclassifier.eval()# Apply the same transformations as we did for the training imagestransformation=transforms.Compose([transforms.ToTensor(),transforms.Normalize(mean=[0.5,0.5,0.5],std=[0.5,0.5,0.5])])# Preprocess the imageimage_tensor=transformation(image).float()# Add an extra batch dimension since pytorch treats all inputs as batchesimage_tensor=image_tensor.unsqueeze_(0)# Turn the input into a Variableinput_features=Variable(image_tensor)# Predict the class of the imageoutput=classifier(input_features)index=output.data.numpy().argmax()returnindex#Now let's try it with a new imagefromrandomimportrandintfromPILimportImageimportos,shutilimgFile=Image.open("classify3.jpg")imgFile=imgFile.resize((resize_x,resize_y),Image.ANTIALIAS)# Display the imageplt.imshow(imgFile)# Call the predction functionindex=predict_image(model,imgFile)print(classes[index])
In [ ]: