anomaly-detection

UCL Final year project that detects anomalies in X-ray images of parcels using a siamese neural network

Quick guide

Run the following script to detect anomalies in a novel image

import torch
import tochvision
import cv2
import numpy as np

from networks import TripletNet, EmbeddingNet
from joblib import load
from generate_vectors import produce_patches

model = TripletNet(EmbeddingNet(torchvision.models.densenet201()))

# Map location should be changed depending on whether CUDA is available or not
model.load_state_dict(torch.load('triplet_densenet201_m2.pth', map_location=torch.device('cpu')))
model.eval()
embedding_net = model.embedding_net

isf = load('isolation_forest.joblib')

embeddings = np.empty((0, 1024))

patches = produce_patches(PATH_TO_IMAGE)

for patch in patches:
    with torch.no_grad():
        vec = embedding_net(patch).cpu().detach().numpy()
    if vec.shape == (1024,):
        vec = np.expand_dims(a, axis=0)
    embeddings = np.append(embeddings, vec, axis=0)
    
y_pred = min(isf.predict(embeddings))
print(y_pred)

triplet_densenet_m2.pth is a file containing the saved weights of the model, which was trained on the X-ray image dataset.

isolation_forest.joblib is a file containing the saved isolation forest, which was trained using the embeddings generated by the model saved in triplet_densenet_m2.pth.

Training

Put all images in a dataset folder. Then run datasets.siamese_train_test_csv or datasets.triplet_train_test_csv to generate the training and testing csv files.

After that, simply run train_siamese.py or train_triplet.py to train the models. When the models are finished training, model.pth will contain the best performing model.

To train the isolation forest, you'll first need to generate the vectors using the model you just trained. Use generate_vectors function to save the generated vectors into files, or just write your own function if you want to use online learning, and then use train_isolation.py to train the isolation forest.