def load_model(model_path):
"""
load model with cpu or gpu mode
"""
local_dir = os.path.abspath("model_file")
client = S3()
client.download(obs_path=model_path, local_dir=local_dir)
file_name = model_path.split('/')[len(model_path.split('/')) - 1]
model = get_segmentation_model(
'danet',
dataset='norm',
backbone='resnet50',
aux=False,
se_loss=False,
norm_layer=torch.nn.BatchNorm2d,
base_size=512,
crop_size=512,
multi_grid=True,
multi_dilation=(4, 8, 16))
try:
if use_cuda:
model = model.to("cuda:0")
checkpoint = torch.load(os.path.join(local_dir, file_name))
else:
checkpoint = torch.load(os.path.join(local_dir, file_name), map_location=torch.device('cpu'))
model.load_state_dict(checkpoint['state_dict'], strict=False)
except Exception as e:
print("error when load model file: ", e)
exit(1)
return model
def save_to_obstiff(self,
obs_path,
feature=None,
no_data_value=None,
merge_method="last"):
"""
Save indexed EOPatches to a complete tiff, and upload to obs.
:param feature: Feature which will be exported
:type feature: (FeatureType, str)
:param obs_path: obs path to save tiff
:type obs_path: str
:param no_data_value: Value of pixels of tiff image with no data in EOPatch
:type no_data_value: int or float
:param merge_method: How to merge overlap EOPatches. "last" mean latter array overwrite former array,
"first" mean former array overwrite latter array.
:type merge_method: str
"""
if not feature:
feature = self.feature
# temp_file = tempfile.TemporaryFile(suffix=".tiff").name
temp_file = tempfile.mkstemp(suffix=".tiff")[1]
try:
self.save_to_tiff(temp_file,
feature=feature,
no_data_value=no_data_value,
merge_method=merge_method)
s3 = S3()
return s3.upload(local_file=temp_file, obs_path=obs_path)
finally:
if os.path.exists(temp_file):
os.remove(temp_file)
def generate_tiff_from_array(meta, array, obs_path):
temp_tiff_file = tempfile.mktemp(suffix=".tiff")
temp_cog_file = tempfile.mktemp(suffix=".tiff")
with rasterio.open(temp_tiff_file, mode='w', driver=meta.get('driver'), width=meta.get('width'),
height=meta.get('height'), transform=meta.get('transform'), crs=meta.get('crs'),
count=meta.get('count'), nodata=meta.get('nodata'), dtype=array.dtype) as dst:
for band in range(0, meta.get('count')):
dst.write_band(band + 1, array[band])
try:
# convert tiff to cog
output_profile = cog_profiles.get("deflate")
output_profile.update({"BLOCKXSIZE": 256, "BLOCKYSIZE": 256})
config = dict(
NUM_THREADS=8,
GDAL_TIFF_INTERNAL_MASK=os.environ.get("GDAL_TIFF_INTERNAL_MASK", True),
GDAL_TIFF_OVR_BLOCKSIZE=str(os.environ.get("GDAL_TIFF_OVR_BLOCKSIZE", 128))
)
cog_translate(temp_tiff_file, temp_cog_file, output_profile, add_mask=True, web_optimized=False,
config=config)
s3 = S3()
return s3.upload(local_file=temp_cog_file, obs_path=obs_path)
except Exception as e:
raise ConvertCogError(e.__str__())
finally:
if os.path.exists(temp_tiff_file):
os.remove(temp_tiff_file)
if os.path.exists(temp_cog_file):
os.remove(temp_cog_file)
def to_obstiff(arr, obs_path, proj="EPSG:4326", spec=None, bands=None, **kwargs):
temp_file = tempfile.mktemp(suffix=".tiff")
temp_cog_file = tempfile.mktemp(suffix=".tiff")
# save to tiff
to_geotiff(arr, path=temp_file, proj=proj, spec=spec, bands=bands, **kwargs)
try:
# convert tiff to cog
output_profile = cog_profiles.get("deflate")
output_profile.update({"BLOCKXSIZE": 256, "BLOCKYSIZE": 256})
config = dict(
NUM_THREADS=8,
GDAL_TIFF_INTERNAL_MASK=os.environ.get("GDAL_TIFF_INTERNAL_MASK", True),
GDAL_TIFF_OVR_BLOCKSIZE=str(os.environ.get("GDAL_TIFF_OVR_BLOCKSIZE", 128))
)
cog_translate(temp_file, temp_cog_file, output_profile, add_mask=True, web_optimized=False, config=config)
s3 = S3()
return s3.upload(local_file=temp_cog_file, obs_path=obs_path)
except Exception as e:
raise ConvertCogError(e.__str__())
finally:
if os.path.exists(temp_file):
os.remove(temp_file)
if os.path.exists(temp_cog_file):
os.remove(temp_cog_file)
def predict_process(output_path, model_path, cat_ids=None, paths=None, aoi=None):
client = S3()
bucket = client.info.get("bucket")
padded_temp_tiff = "obs://%s/temp.tif" % bucket
try:
print("start check environment variable...")
obs_env_check()
print("start load model...")
model = load_model(model_path)
print("start load data...")
cat_ids = split2list(cat_ids)
paths = split2list(paths)
if aoi is not None:
aoi = split2list(aoi)
aoi = [float(i) for i in aoi]
patch_set = load_data(cat_ids=cat_ids, paths=paths, padded_temp_tiff=padded_temp_tiff, aoi=aoi)
print("start predict...")
predict_model(model=model, patch_set=patch_set)
print("write to obs...")
patch_set.save_to_obstiff(output_path, (FeatureType.DATA, 'MASK'), no_data_value=0, padding=20)
print("predict model done")
finally:
client.delete(padded_temp_tiff)
def generate_padded_tiff(meta, array, pad_width, obs_path, pad_mode='constant'):
temp_tiff_file = tempfile.mktemp(suffix=".tiff")
temp_cog_file = tempfile.mktemp(suffix=".tiff")
if len(array.shape) != len(pad_width):
raise ValueError('Pad width is invalid')
new_width = meta.get('width') + pad_width[1][0] + pad_width[1][1]
new_height = meta.get('height') + pad_width[2][0] + pad_width[2][1]
padded_array, new_transform = pad(array, meta.get('transform'), pad_width=pad_width, mode=pad_mode)
# write padded tiff
with rasterio.open(temp_tiff_file, mode='w', driver=meta.get('driver'), width=new_width,
height=new_height, transform=new_transform, crs=meta.get('crs'),
count=meta.get('count'), nodata=meta.get('nodata'), dtype=meta.get('dtype')) as dst:
for band in range(0, meta.get('count')):
dst.write_band(band + 1, padded_array[band])
try:
# convert tiff to cog
output_profile = cog_profiles.get("deflate")
output_profile.update({"BLOCKXSIZE": 256, "BLOCKYSIZE": 256})
config = dict(
NUM_THREADS=8,
GDAL_TIFF_INTERNAL_MASK=os.environ.get("GDAL_TIFF_INTERNAL_MASK", True),
GDAL_TIFF_OVR_BLOCKSIZE=str(os.environ.get("GDAL_TIFF_OVR_BLOCKSIZE", 128))
)
cog_translate(temp_tiff_file, temp_cog_file, output_profile, nodata=0, add_mask=True, web_optimized=False,
config=config)
s3 = S3()
return s3.upload(local_file=temp_cog_file, obs_path=obs_path)
except Exception as e:
raise ConvertCogError(e.__str__())
finally:
if os.path.exists(temp_tiff_file):
os.remove(temp_tiff_file)
if os.path.exists(temp_cog_file):
os.remove(temp_cog_file)
def __init__(self, args, pre_train_model):
self.args = args
args.log_name = str(args.checkname)
self.logger = utils.create_logger(args.log_root, args.log_name)
# data transforms
input_transform = transform.Compose([
transform.ToTensor(),
transform.Normalize([.485, .456, .406], [.229, .224, .225])
])
# dataset
data_kwargs = {
'transform': input_transform,
'base_size': args.base_size,
'crop_size': args.crop_size,
'logger': self.logger,
'scale': args.scale
}
print('start copy data')
client = S3()
client.download(pre_train_model, './pretrain_models')
print('copy success')
print('num gpus ' + str(torch.cuda.device_count()))
trainset = get_segmentation_dataset(args.dataset,
root=args.data_folder,
split='train',
mode='train',
**data_kwargs)
testset = get_segmentation_dataset(args.dataset,
root=args.data_folder,
split='val',
mode='val',
**data_kwargs)
# dataloader
kwargs = {'num_workers': args.workers, 'pin_memory': True} \
if args.cuda else {}
self.trainloader = data.DataLoader(trainset,
batch_size=args.batch_size,
drop_last=True,
shuffle=True,
**kwargs)
self.valloader = data.DataLoader(testset,
batch_size=args.batch_size,
drop_last=False,
shuffle=False,
**kwargs)
self.nclass = trainset.num_class
# model
model = get_segmentation_model(args.model,
dataset=args.dataset,
backbone=args.backbone,
aux=args.aux,
se_loss=args.se_loss,
norm_layer=torch.nn.BatchNorm2d,
base_size=args.base_size,
crop_size=args.crop_size,
multi_grid=args.multi_grid,
multi_dilation=args.multi_dilation)
self.logger.info(model)
# optimizer using different LR
params_list = [
{
'params': model.pretrained.parameters(),
'lr': args.lr
},
]
if hasattr(model, 'head'):
params_list.append({
'params': model.head.parameters(),
'lr': args.lr * 10
})
if hasattr(model, 'auxlayer'):
params_list.append({
'params': model.auxlayer.parameters(),
'lr': args.lr * 10
})
optimizer = torch.optim.SGD(params_list,
lr=args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
self.criterion = SegmentationMultiLosses(nclass=self.nclass)
self.model, self.optimizer = model, optimizer
# using cuda
if args.cuda:
self.model = DataParallelModel(self.model).cuda()
self.criterion = DataParallelCriterion(self.criterion).cuda()
# finetune from a trained model
if args.ft:
args.start_epoch = 0
checkpoint = torch.load(args.ft_resume)
if args.cuda:
self.model.module.load_state_dict(checkpoint['state_dict'],
strict=False)
else:
self.model.load_state_dict(checkpoint['state_dict'],
strict=False)
self.logger.info("=> loaded checkpoint '{}' (epoch {})".format(
args.ft_resume, checkpoint['epoch']))
# resuming checkpoint
if args.resume:
if not os.path.isfile(args.resume):
raise RuntimeError("=> no checkpoint found at '{}'".format(
args.resume))
checkpoint = torch.load(args.resume)
args.start_epoch = checkpoint['epoch']
if args.cuda:
self.model.module.load_state_dict(checkpoint['state_dict'])
else:
self.model.load_state_dict(checkpoint['state_dict'])
if not args.ft:
self.optimizer.load_state_dict(checkpoint['optimizer'])
self.best_pred = checkpoint['best_pred']
self.logger.info("=> loaded checkpoint '{}' (epoch {})".format(
args.resume, checkpoint['epoch']))
# lr scheduler
self.scheduler = utils.LR_Scheduler(args.lr_scheduler,
args.lr,
args.epochs,
len(self.trainloader),
logger=self.logger,
lr_step=args.lr_step)
self.best_pred = 0.0
def predict_process_pv(model_path, img_path, img_obs_path, output_path):
local_dir = os.path.abspath("model_file")
client = S3()
client.download(obs_path=model_path, local_dir=local_dir)
file_name = model_path.split('/')[len(model_path.split('/')) - 1]
model = get_segmentation_model('danet',
dataset='norm',
backbone='resnet50',
aux=False,
se_loss=False,
norm_layer=torch.nn.BatchNorm2d,
base_size=512,
crop_size=512,
multi_grid=True,
multi_dilation=(4, 8, 16))
try:
if use_cuda:
model = model.to("cuda:0")
checkpoint = torch.load(os.path.join(local_dir, file_name))
else:
checkpoint = torch.load(os.path.join(local_dir, file_name),
map_location=torch.device('cpu'))
model.load_state_dict(checkpoint['state_dict'], strict=False)
except Exception as e:
print("error when load model file: ", e)
exit(1)
large_img = cv2.imread(img_path, 4)
large_img = cv2.copyMakeBorder(large_img,
pad,
pad,
pad,
pad,
cv2.BORDER_CONSTANT,
value=(0, 0, 0))
large_img = cv2.cvtColor(large_img, cv2.COLOR_BGR2RGB)
data_list, offset_list = split_data(large_img)
large_bin_pred = np.zeros(large_img.shape[:2])
print(len(data_list))
for idx, input_data in enumerate(data_list):
if use_cuda:
input_data = input_data.cuda(0)
output = model.evaluate(input_data)
pre = torch.max(output, 1)[1].cpu().numpy() + 1
for sub_idx in range(BATCH_SIZE):
if (idx * BATCH_SIZE + sub_idx) >= len(offset_list):
break
offset_h = offset_list[idx * BATCH_SIZE + sub_idx][0]
offset_w = offset_list[idx * BATCH_SIZE + sub_idx][1]
large_bin_pred[offset_h + pad:offset_h + sub_h - pad, offset_w +
pad:offset_w + sub_w - pad] = pre[sub_idx][pad:-pad,
pad:-pad]
result = large_bin_pred[pad:-pad, pad:-pad]
origin_img = OBSImage(img_obs_path)
mask_img = np.max(origin_img.read().transpose(1, 2, 0)[:, :, :3], axis=2)
mask_img[mask_img > 0] = 1
result = result * mask_img
factory = TiffFactory()
factory.generate_tiff_from_array(origin_img.get_image_meta(), result,
output_path)