BuildingExtraction / Tools /GetImgMeanStd.py
KyanChen's picture
add model
ab01e4a
import os
import pandas as pd
from skimage import io
import numpy as np
import json
import tqdm
# R, G, B
class GetImgMeanStd:
def __init__(self, data_file):
assert os.path.exists(data_file), 'train.csv dose not exist!'
self.data_info = pd.read_csv(data_file, index_col=0)
self.save_path_mean_std_info = 'generate_dep_info'
self.mean = None
self.std = None
def get_img_mean_std(self):
means = []
stds = []
bar = tqdm.tqdm(total=len(self.data_info))
for row in self.data_info.iterrows():
bar.update(1)
img_name = row[1]['img']
# print(img_name)
img = io.imread(img_name)
img = img / 255.
assert img is not None, img_name + 'is not valid'
# height*width*channels, axis=0 is the first dim
mean = np.mean(np.mean(img, axis=0), axis=0)
means.append(mean)
std = np.std(np.std(img, axis=0), axis=0)
stds.append(std)
bar.close()
self.mean = np.mean(np.array(means), axis=0).tolist()
self.std = np.mean(np.array(stds), axis=0).tolist()
return {'mean': self.mean, 'std': self.std}
def write_mean_std_information(self):
info = self.get_img_mean_std()
writer = os.path.join(self.save_path_mean_std_info, 'mean_std_info_test.json')
with open(writer, 'w') as f_writer:
json.dump(info, f_writer)
print('\'PRIOR_MEAN\': %s\n\'PRIOR_STD\': %s\n' % (info['mean'], info['std']))
if __name__ == '__main__':
data_file = r'generate_dep_info/train_data.csv'
getImgMeanStd = GetImgMeanStd(data_file)
getImgMeanStd.write_mean_std_information()