Skip to content

Added notebooks for CNN-RNN and Spacenet6 #4

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
popovstefan wants to merge 1 commit into
base: master
Choose a base branch
from
Open

Added notebooks for CNN-RNN and Spacenet6 #4

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion aitlas/datasets/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -31,7 +31,7 @@
from .amazon_rainforest import AmazonRainforestDataset
from .airs import AIRSDataset
from .camvid import CamVidDataset
# from .spacenet6 import SpaceNet6Dataset
from .spacenet6 import SpaceNet6Dataset
from .massachusetts_buildings import MassachusettsBuildingsDataset
from .massachusetts_roads import MassachusettsRoadsDataset
from .inria import InriaDataset
Expand Down
180 changes: 59 additions & 121 deletions aitlas/datasets/spacenet6.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,20 +12,19 @@
from multiprocessing import Pool

import cv2
import gdal
from osgeo import gdal
import numpy as np
import pandas as pd
import torch
from shapely.wkt import loads
from skimage import io, measure
from skimage.morphology import dilation, erosion, square
from skimage.morphology import square, erosion, dilation
from skimage.segmentation import watershed
from tqdm import tqdm

from ..base import BaseDataset
from ..datasets.schemas import SpaceNet6DatasetSchema
from ..utils import parse_img_id

from aitlas.base import BaseDataset
from aitlas.datasets.schemas import SpaceNet6DatasetSchema
from aitlas.utils.utils import parse_img_id

# Ignore the "low-contrast" warnings
warnings.filterwarnings("ignore")
Expand All @@ -44,9 +43,7 @@ def integer_coordinates(x):
return image_mask


def process_image(
image_path, segmentation_directory, edge_width, contact_width, gt_buildings_csv
):
def process_image(image_path, segmentation_directory, edge_width, contact_width, gt_buildings_csv):
"""
Creates and saves the target (ground-truth) segmentation mask for the input image.

Expand All @@ -65,11 +62,10 @@ def process_image(
"""
gt_buildings = pd.read_csv(gt_buildings_csv)
image_name = os.path.basename(image_path)
values = gt_buildings[
(gt_buildings["ImageId"] == "_".join(image_name.split("_")[-4:])[:-4])
][["TileBuildingId", "PolygonWKT_Pix", "Mean_Building_Height"]].values
labels = np.zeros((900, 900), dtype="uint16")
heights = np.zeros((900, 900), dtype="float")
values = gt_buildings[(gt_buildings["ImageId"] == "_".join(image_name.split("_")[-4:])[:-4])][
["TileBuildingId", "PolygonWKT_Pix", "Mean_Building_Height"]].values
labels = np.zeros((900, 900), dtype='uint16')
heights = np.zeros((900, 900), dtype='float')
cur_lbl = 0
for i in range(values.shape[0]):
poly = loads(values[i, 1])
Expand Down Expand Up @@ -103,31 +99,24 @@ def process_image(
sz = 3
else:
sz = 1
unique = np.unique(
labels[
max(0, y0 - sz) : min(labels.shape[0], y0 + sz + 1),
max(0, x0 - sz) : min(labels.shape[1], x0 + sz + 1),
]
)
unique = np.unique(labels[max(0, y0 - sz):min(labels.shape[0], y0 + sz + 1),
max(0, x0 - sz):min(labels.shape[1], x0 + sz + 1)])
if len(unique[unique > 0]) > 1:
contact_msk[y0, x0] = True
msk = np.stack(
(255 * footprint_msk, 255 * border_msk, 255 * contact_msk)
).astype("uint8")
msk = np.stack((255 * footprint_msk, 255 * border_msk, 255 * contact_msk)).astype("uint8")
msk = np.rollaxis(msk, 0, 3)
io.imsave(os.path.join(segmentation_directory, image_name), msk)


class SpaceNet6Dataset(BaseDataset):
"""SpaceNet6 dataset."""

schema = SpaceNet6DatasetSchema

def __init__(self, config):
BaseDataset.__init__(self, config)
self.image_paths = list()
self.mask_paths = list()
self.orients = pd.read_csv(config.orients, index_col=0)
self.orients = pd.read_csv(config["orients"], index_col=0)
self.orients["val"] = list(range(len(self.orients.index)))

def __getitem__(self, index):
Expand All @@ -151,34 +140,35 @@ def __getitem__(self, index):
m = np.where((image.sum(axis=2) > 0).any(0))
x_min, x_max = np.amin(m), np.amax(m) + 1
# Remove black parts
image = image[y_min:y_max, x_min:x_max]
image = image[y_min: y_max, x_min: x_max]
# Apply transformations, (should be available only for training data)
if self.config.transforms:
# Get mask path
mask_path = self.mask_paths[index]
# Read mask
mask = io.imread(mask_path)
# Remove black parts
mask = mask[y_min:y_max, x_min:x_max]
image, mask = self.transform({"image": image, "mask": mask})
mask = mask[y_min: y_max, x_min: x_max]
image, mask = self.transform({
"image": image,
"mask": mask
})
# Extract direction, strip and coordinates from image
direction, strip, coordinate = parse_img_id(image_path, self.orients)
if direction.item():
image = np.fliplr(np.flipud(image))
if self.config.transforms:
mask = np.fliplr(np.flipud(mask))
image = (
image - np.array([28.62501827, 36.09922463, 33.84483687, 26.21196667])
) / np.array([8.41487376, 8.26645475, 8.32328472, 8.63668993])
image = (image - np.array([28.62501827, 36.09922463, 33.84483687, 26.21196667])) / np.array(
[8.41487376, 8.26645475, 8.32328472, 8.63668993])
# Transpose image
image = torch.from_numpy(image.transpose((2, 0, 1)).copy()).float()
# Reorder bands
image = image[[0, 3, 1, 2]]
if self.config.transforms:
weights = np.ones_like(mask[:, :, :1], dtype=float)
region_labels, region_count = measure.label(
mask[:, :, 0], background=0, connectivity=1, return_num=True
)
region_labels, region_count = measure.label(mask[:, :, 0], background=0, connectivity=1,
return_num=True)
region_properties = measure.regionprops(region_labels)
for bl in range(region_count):
weights[region_labels == bl + 1] = 1024.0 / region_properties[bl].area
Expand All @@ -199,7 +189,7 @@ def __getitem__(self, index):
"ymin": y_min,
"xmin": x_min,
"b_count": region_count,
"weights": weights,
"weights": weights
}

def __len__(self):
Expand All @@ -213,34 +203,16 @@ def load_directory(self):
def load_other_folds(self, fold):
"""Loads all images (and masks) except the ones from this fold."""
df = pd.read_csv(self.config.folds_path)
self.image_paths = [
os.path.join(
self.config.root_directory, "SAR-Intensity", os.path.basename(x)
)
for x in df[
np.logical_or(
df["fold"] > (fold % 10) + 1, df["fold"] < (fold % 10) - 1
)
]["sar"].values
]
self.mask_paths = [
os.path.join(self.config.segmentation_directory, os.path.basename(x))
for x in df[
np.logical_or(
df["fold"] > (fold % 10) + 1, df["fold"] < (fold % 10) - 1
)
]["segm"].values
]
self.image_paths = [os.path.join(self.config.root_directory, "SAR-Intensity", os.path.basename(x)) for x in
df[np.logical_or(df["fold"] > (fold % 10) + 1, df["fold"] < (fold % 10) - 1)]["sar"].values]
self.mask_paths = [os.path.join(self.config.segmentation_directory, os.path.basename(x)) for x in
df[np.logical_or(df["fold"] > (fold % 10) + 1, df["fold"] < (fold % 10) - 1)]["segm"].values]

def load_fold(self, fold):
"""Loads the images from this fold."""
df = pd.read_csv(self.config.folds_path)
self.image_paths = [
os.path.join(
self.config.root_directory, "SAR-Intensity", os.path.basename(x)
)
for x in df[df["fold"] == (fold % 10)]["sar"].values
]
self.image_paths = [os.path.join(self.config.root_directory, "SAR-Intensity", os.path.basename(x)) for x in
df[df["fold"] == (fold % 10)]["sar"].values]
self.mask_paths = None

def labels(self):
Expand All @@ -261,87 +233,53 @@ def prepare(self):
os.makedirs(self.config.segmentation_directory, exist_ok=True)
os.makedirs(self.config.folds_dir, exist_ok=True)
# Path to the ground-truth buildings csv file
gt_buildings_csv_filepath = os.path.join(
self.config.root_directory,
"SummaryData/SN6_Train_AOI_11_Rotterdam_Buildings.csv",
)
gt_buildings_csv_filepath = os.path.join(self.config.root_directory,
"SummaryData/SN6_Train_AOI_11_Rotterdam_Buildings.csv")
# Read gt building csv file
gt_buildings = pd.read_csv(gt_buildings_csv_filepath)
# Walk the raw data directory with the SAR images and save the filenames in it
sar_image_paths = glob.glob(
os.path.join(self.config.root_directory, "SAR-Intensity", "*.tif")
)
sar_image_paths = glob.glob(os.path.join(self.config.root_directory, "SAR-Intensity", "*.tif"))
# Process each SAR image
with Pool(self.config.num_threads) as pool:
for _ in tqdm(
pool.imap_unordered(
partial(
process_image,
segmentation_directory=self.config.segmentation_directory,
edge_width=self.config.edge_width,
contact_width=self.config.contact_width,
gt_buildings_csv=gt_buildings_csv_filepath,
),
sar_image_paths,
)
):
for _ in tqdm(pool.imap_unordered(
partial(process_image,
segmentation_directory=self.config.segmentation_directory,
edge_width=self.config.edge_width,
contact_width=self.config.contact_width,
gt_buildings_csv=gt_buildings_csv_filepath),
sar_image_paths)):
pass
orientations = pd.read_csv(
filepath_or_buffer=self.config.orients,
sep=" ",
index_col=0,
names=["strip", "direction"],
header=None,
)
df_fold = pd.DataFrame(
columns=["ImageId", "sar", "segm", "rotation", "x", "y", "fold"]
)
orientations = pd.read_csv(filepath_or_buffer=self.config.orients, sep=' ', index_col=0,
names=["strip", "direction"], header=None)
df_fold = pd.DataFrame(columns=["ImageId", "sar", "segm", "rotation", "x", "y", "fold"])
l_edge = 591640
r_edge = 596160
orientations["sum_y"] = 0.0
orientations["ctr_y"] = 0.0
for sar_path in tqdm(sar_image_paths):
image_id = "_".join(
os.path.splitext(os.path.basename(sar_path))[0].split("_")[-4:]
)
image_id = "_".join(os.path.splitext(os.path.basename(sar_path))[0].split("_")[-4:])
strip_name = "_".join(image_id.split("_")[-4:-2])
rotation = orientations.loc[strip_name]["direction"].squeeze()
tr = gdal.Open(sar_path).GetGeoTransform()
orientations.loc[strip_name, "sum_y"] += tr[3]
orientations.loc[strip_name, "ctr_y"] += 1
fold_no = min(
self.config.num_folds - 1,
max(
0,
math.floor(
(tr[0] - l_edge) / (r_edge - l_edge) * self.config.num_folds
),
),
)
segmentation_path = os.path.join(
self.config.segmentation_directory, os.path.basename(sar_path)
)
df_fold = df_fold.append(
{
"ImageId": image_id,
"sar": sar_path,
"segm": segmentation_path,
"rotation": rotation,
"x": tr[0],
"y": tr[3],
"fold": fold_no,
},
ignore_index=True,
)
fold_no = min(self.config.num_folds - 1,
max(0, math.floor((tr[0] - l_edge) / (r_edge - l_edge) * self.config.num_folds)))
segmentation_path = os.path.join(self.config.segmentation_directory, os.path.basename(sar_path))
df_fold = df_fold.append({
"ImageId": image_id,
"sar": sar_path,
"segm": segmentation_path,
"rotation": rotation,
"x": tr[0],
"y": tr[3],
"fold": fold_no
}, ignore_index=True)
df_fold.to_csv(os.path.join(self.config.folds_dir, "folds.csv"), index=False)
for i in range(self.config.num_folds):
img_ids = df_fold[df_fold["fold"] == i]["ImageId"].values
gt_buildings[gt_buildings.ImageId.isin(img_ids)].to_csv(
os.path.join(self.config.folds_dir, "gt_fold{}.csv").format(i),
index=False,
)
os.path.join(self.config.folds_dir, "gt_fold{}.csv").format(i), index=False)
orientations["mean_y"] = orientations["sum_y"] / orientations["ctr_y"]
orientations["coord_y"] = (
(orientations["mean_y"] - 5746153.106161971) / 11000
) + 0.2
orientations["coord_y"] = (((orientations["mean_y"] - 5746153.106161971) / 11000) + 0.2)
orientations.to_csv(self.config.orients_output, index=True)
2 changes: 1 addition & 1 deletion aitlas/models/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,5 +18,5 @@
EfficientNetB0, EfficientNetB0MultiLabel, EfficientNetB4, EfficientNetB4MultiLabel,
EfficientNetB7, EfficientNetB7MultiLabel
)
# from .unet_efficientnet import UNetEfficientNet
from .unet_efficientnet import UNetEfficientNet
from .vgg import VGG16, VGG19, VGG16MultiLabel, VGG19MultiLabel
Loading