text_plain_output_1.png
from torch.nn import CrossEntropyLoss

from torch.utils.data import random_split, DataLoader, Dataset

from torchvision.io import read_image, ImageReadMode

from torchvision.models.segmentation.deeplabv3 import DeepLabHead

from torchvision.models.segmentation.fcn import FCNHead

from tqdm import tqdm

from tqdm.notebook import tqdm

import os

import torch

import torch.nn.functional as F

import torch.nn.functional as F

import torchvision

import torchvision.transforms as transforms

image_path = '/content/dataset/semantic_drone_dataset/original_images'
mask_path = '/content/dataset/semantic_drone_dataset/label_images_semantic'

length = len(os.listdir(image_path))

class DroneDataset(Dataset):
    def __init__(self, imgs_dir, masks_dir, count, is_val=False):
        self.imgs_dir = imgs_dir
        self.masks_dir = masks_dir
        imgs_paths = os.listdir(self.imgs_dir)
        imgs_paths.sort()
        mask_paths = os.listdir(self.masks_dir)
        mask_paths.sort()
        self.is_val = is_val
        if not is_val:
            self.imgs_paths = imgs_paths[:count]
            self.mask_paths = mask_paths[:count]
        else:
            self.imgs_paths = imgs_paths[-count:]
            self.mask_paths = mask_paths[-count:]
    def __len__(self):
        return len(self.imgs_paths)
    def __getitem__(self, idx):
        img = read_image(os.path.join(self.imgs_dir, self.imgs_paths[idx]), ImageReadMode.RGB)
        mask = read_image(os.path.join(self.masks_dir, self.mask_paths[idx]), ImageReadMode.GRAY)
        return (img, mask)

torchvision.models.segmentation.DeepLabV3_ResNet50_Weights.COCO_WITH_VOC_LABELS_V1.transforms()

def img_transform(img, mask, is_val=False, size=520):
    img = img.to(device)
    mask = mask.to(device)
    img = img.float() / 255.0
    if not is_val:
        trans_img = torch.nn.Sequential(transforms.Resize([size, size]), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), transforms.RandomAutocontrast(p=0.2))
    else:
        trans_img = trans_img = torch.nn.Sequential(transforms.Resize([size, size]), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]))
    trans_mask = torch.nn.Sequential(transforms.Resize([size, size]))
    trans_img.requires_grad_(False)
    trans_mask.requires_grad_(False)
    trans_img = trans_img.to(device)
    trans_mask = trans_mask.to(device)
    img = trans_img(img)
    mask = trans_mask(mask)
    return (img, mask.squeeze(1).long())

train_dataset_len = int(length * 0.7)
val_dataset_len = length - train_dataset_len
train_dataset = DroneDataset(image_path, mask_path, train_dataset_len)
val_dataset = DroneDataset(image_path, mask_path, val_dataset_len, is_val=True)

batch_size = 4
train_loader = DataLoader(train_dataset, batch_size, shuffle=True, num_workers=2)
val_loader = DataLoader(val_dataset, batch_size, shuffle=False, num_workers=2)

model = torchvision.models.segmentation.deeplabv3_resnet50(weights=torchvision.models.segmentation.DeepLabV3_ResNet50_Weights.DEFAULT, progress=True)

from torchvision.models.segmentation.deeplabv3 import DeepLabHead
from torchvision.models.segmentation.fcn import FCNHead
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
model.classifier = DeepLabHead(2048, 23)
model.aux_classifier = FCNHead(1024, 23)
model = model.to(device)

from torch.nn import CrossEntropyLoss
import torch.nn.functional as F
loss = CrossEntropyLoss().to(device)
learning_rate = 0.01
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)

def pixel_accuracy(mask, output):
    output_softmax = F.softmax(output, dim=1)
    output_argmax = torch.argmax(output_softmax, dim=1)
    bool_tensor = torch.flatten(mask) == torch.flatten(output_argmax)
    return torch.sum(bool_tensor) / torch.numel(bool_tensor)

from tqdm import tqdm
epoch_count = 30
train_losses = []
val_losses = []
train_accs = []
val_accs = []
es_steps = 3
count_steps = 0
train_len = len(train_loader)
val_len = len(val_loader)
print(train_len)
print(val_len)
best_score = 10000000000.0
for epoch in range(epoch_count):
    if count_steps >= es_steps:
        print('Early stopping!')
        break
    train_loss_sum = 0
    train_pixel_acc = 0
    model.train()
    for img_batch, mask_batch in tqdm(train_loader):
        img_batch = img_batch.to(device, non_blocking=True)
        mask_batch = mask_batch.to(device, non_blocking=True)
        img_batch, mask_batch = img_transform(img_batch, mask_batch, is_val=False)
        optimizer.zero_grad()
        output_batch = model(img_batch)
        loss_value = loss(output_batch['out'], mask_batch)
        train_pixel_acc += pixel_accuracy(mask_batch, output_batch['out']).detach()
        train_loss_sum += loss_value.detach()
        loss_value.backward()
        optimizer.step()
        del output_batch
    train_loss = train_loss_sum / train_len
    train_acc = train_pixel_acc / train_len
    train_losses.append(train_loss)
    train_accs.append(train_acc)
    print(f'Epoch {epoch} / {epoch_count} | train loss = {train_loss} | train acc = {train_acc}')
    model.eval()
    val_loss_sum = 0
    val_pixel_acc = 0
    for img_batch, mask_batch in tqdm(val_loader):
        img_batch = img_batch.to(device, non_blocking=True)
        mask_batch = mask_batch.to(device, non_blocking=True)
        img_batch, mask_batch = img_transform(img_batch, mask_batch, is_val=True)
        output_batch = model(img_batch)
        loss_value = loss(output_batch['out'], mask_batch)
        val_loss_sum = val_loss_sum + loss_value.detach()
        val_pixel_acc = val_pixel_acc + pixel_accuracy(mask_batch, output_batch['out']).detach()
        del output_batch
    val_loss = val_loss_sum / val_len
    val_acc = val_pixel_acc / val_len
    val_losses.append(val_loss)
    val_accs.append(val_acc)
    print(f'Epoch {epoch} / {epoch_count} | val loss = {val_loss} | val acc = {val_acc}')
    if val_loss < best_score:
        best_score = val_loss
        count_steps = 0
        torch.save(model, 'best_model.pt')
    else:
        count_steps += 1