64 lines
1.5 KiB
Python
64 lines
1.5 KiB
Python
import os
|
|
import numpy as np
|
|
import matplotlib.pyplot as plt
|
|
|
|
import cv2 as cv
|
|
import rasterio as rio
|
|
|
|
RASTERS_DIR = './data/data/train_features'
|
|
IMAGES_DIR = './images/'
|
|
FC_DIR = "fc"
|
|
|
|
|
|
def load_img(path, expand_dim=False):
|
|
img = cv.imread(path)
|
|
img = img / 255.0
|
|
if expand_dim:
|
|
img = np.expand_dims(img, axis=0)
|
|
|
|
return img
|
|
|
|
|
|
def create_folder(name, path):
|
|
n = os.path.join(path, name)
|
|
if not os.path.exists(n):
|
|
os.makedirs(n)
|
|
return n
|
|
|
|
|
|
def scale(band):
|
|
return band/np.max(band)
|
|
|
|
|
|
def convert_raster_to_image(rasters_dir,
|
|
rgb_path=None,
|
|
false_color_path=None):
|
|
|
|
b2 = scale(rio.open(rasters_dir+'/B02.tif').read().reshape(512, 512, 1))
|
|
b3 = scale(rio.open(rasters_dir+'/B03.tif').read().reshape(512, 512, 1))
|
|
b4 = scale(rio.open(rasters_dir+'/B04.tif').read().reshape(512, 512, 1))
|
|
b8 = scale(rio.open(rasters_dir+'/B08.tif').read().reshape(512, 512, 1))
|
|
|
|
file_name = rasters_dir.split(os.sep)[-1]
|
|
|
|
rgb = np.dstack([b4, b3, b2])
|
|
|
|
plt.imsave(fname=rgb_path + f'/{file_name}.jpeg',
|
|
arr=rgb)
|
|
|
|
fc = np.dstack([b8, b3, b2])
|
|
|
|
plt.imsave(fname=false_color_path + f'/{file_name}.jpeg',
|
|
arr=fc)
|
|
|
|
|
|
def transform_photo(raster_dir):
|
|
dp = create_folder(raster_dir, IMAGES_DIR)
|
|
fc = create_folder(FC_DIR, dp)
|
|
convert_raster_to_image(os.path.join(RASTERS_DIR, raster_dir), dp, fc)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
for d in os.listdir(RASTERS_DIR):
|
|
transform_photo(d)
|