NTU-CAE-Segmentation

My project on image segmentation for tile spalling detection at NTUCE CAE Division / NCREE Internship.

The aim is to provide a image segmentation mask for tile spalling of building exterior. We use U-Net, a Deep Learning-based model architecture for the job.

Files

• myunet.py: The code for the altered model architecture. It is modified from segmentation-models-pytorch.
• auto_evaluate.py: A script to automate the evaluation of trained models using testing data.
• Unet-efficientnet-b6-CrossEntropyLoss-4.pt: A sample of the model trained by us. You can check its prediction result by the following sample code:

def predict_image_mask_miou(model, image, mask, mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]):
    model.eval()
    t = T.Compose([T.ToTensor(), T.Normalize(mean, std)])
    image = t(image)
    model.to(device); image=image.to(device)
    mask = mask.to(device)
    with torch.no_grad():
        image = image.unsqueeze(0)
        mask = mask.unsqueeze(0)

        output = model(image)
        masked = torch.argmax(output, dim=1)
        masked = masked.cpu().squeeze(0)
    return masked

Some base code here has originated from tanishqgautam/Drone-Image-Semantic-Segmentation.

Poster