利用Python将原数据压缩到一个pickle 文件夹里

来源:互联网 发布:看股票最好的软件 编辑:程序博客网 时间:2024/05/22 03:26
话不多说。直接上代码。# Process raw data and save them into pickle file.import osimport numpy as npfrom PIL import Imagefrom PIL import ImageOpsfrom scipy import miscimport scipy.iofrom skimage import ioimport cv2import sysimport cPickle as pickleimport globimport randomfrom tqdm import tqdmfrom eliaLib import dataRepresentationfrom constants import *listImgFiles = [k.split('/')[-1].split('.')[0] for k in glob.glob(os.path.join('/home/allen/reface-salgan-master/faceRec/rawface', '*'))]listTestImage = [k.split('/')[-1].split('.')[0] for k in glob.glob(os.path.join('/home/allen/reface-salgan-master/faceRec/rawfacetest', '*'))]# LOAD DATA## TraintrainData = []for currFile in tqdm(listImgFiles):    trainData.append(dataRepresentation.Target(os.path.join('/home/allen/reface-salgan-master/faceRec/lowface', currFile + '.jpg'),                                               os.path.join('/home/allen/reface-salgan-master/faceRec/rawface', currFile + '.jpg'),                                               dataRepresentation.LoadState.loaded, dataRepresentation.InputType.image,                                               dataRepresentation.LoadState.loaded, dataRepresentation.InputType.image,                                               ))with open(os.path.join('/home/allen/reface-salgan-master/faceRec', 'trainData.pickle'), 'wb') as f:    pickle.dump(trainData, f)# testtestData = []for currFile in tqdm(listTestImage):    testData.append(dataRepresentation.Target(os.path.join('/home/allen/reface-salgan-master/faceRec/lowfacetest', currFile + '.jpg'),                                              os.path.join('/home/allen/reface-salgan-master/faceRec/rawfacetest', currFile + '.jpg'),                                              dataRepresentation.LoadState.loaded, dataRepresentation.InputType.image,                                              dataRepresentation.LoadState.loaded, dataRepresentation.InputType.image,                                              ))with open(os.path.join('/home/allen/reface-salgan-master/faceRec', 'testData.pickle'), 'wb') as f:    pickle.dump(testData, f)

另一个.py文件

import cv2import numpy as npfrom enum import Enumimport scipy.ioclass InputType(Enum):    image = 0    imageGrayscale = 1    saliencyMapMatlab = 2    fixationMapMatlab = 3    empty = 100class LoadState(Enum):    unloaded = 0    loaded = 1    loadedCompressed = 2    error = 100###############################################################################################class ImageContainer:    def __init__(self, filePath, lowImageType, state=LoadState.unloaded):        self.filePath = filePath        self.state = state        self.lowImageType = lowImageType        if self.state == LoadState.unloaded:            self.data = None        elif self.state == LoadState.loaded:            self.load()        elif self.state == LoadState.loadedCompressed:            self.loadCompressed()        else:            raise Exception('Unknown state when loading image')    def load(self):        if self.lowImageType == InputType.image:            self.data = cv2.cvtColor(cv2.imread(self.filePath, cv2.IMREAD_COLOR), cv2.COLOR_BGR2RGB)            self.state = LoadState.loaded        if self.lowImageType == InputType.imageGrayscale:            self.data = cv2.cvtColor(cv2.imread(self.filePath, cv2.IMREAD_COLOR), cv2.COLOR_BGR2GRAY)            self.state = LoadState.loaded        elif self.lowImageType == InputType.saliencyMapMatlab:            self.data = (scipy.io.loadmat(self.filePath)['I'] * 255).astype(np.uint8)            self.state = LoadState.loaded        elif self.lowImageType == InputType.fixationMapMatlab:            self.data = (scipy.io.loadmat(self.filePath)['I']).nonzero()            self.state = LoadState.loaded        elif self.lowImageType == InputType.empty:            self.data = None    def loadCompressed(self):        if self.lowImageType == InputType.image:            with open(self.filePath, 'rb') as f:                data = f.read()            self.data = np.fromstring(data, np.uint8)            self.state = LoadState.loadedCompressed        elif self.lowImageType == InputType.saliencyMapMatlab:            self.state = LoadState.error            raise Exception('Saliency maps do no have compressed handlind method enabled')        elif self.lowImageType == InputType.empty:            self.state = LoadState.error            raise Exception('Empty images do no have compressed handlind method enabled')    def getImage(self):        if self.lowImageType == InputType.image:            if self.state == LoadState.unloaded:                return cv2.cvtColor(cv2.imread(self.filePath, cv2.IMREAD_COLOR), cv2.COLOR_BGR2RGB)            elif self.state == LoadState.loaded:                return self.data            elif self.state == LoadState.loadedCompressed:                return cv2.cvtColor(cv2.imdecode(self.data, cv2.IMREAD_COLOR), cv2.COLOR_BGR2RGB)        elif self.lowImageType == InputType.imageGrayscale:            if self.state == LoadState.unloaded:                return cv2.cvtColor(cv2.imread(self.filePath, cv2.IMREAD_COLOR), cv2.COLOR_BGR2GRAY)            elif self.state == LoadState.loaded:                return self.data            elif self.state == LoadState.loadedCompressed:                return cv2.cvtColor(cv2.imdecode(self.data, cv2.IMREAD_COLOR), cv2.COLOR_BGR2GRAY)        elif self.lowImageType == InputType.saliencyMapMatlab:            if self.state == LoadState.unloaded:                return (scipy.io.loadmat(self.filePath)['I'] * 255).astype(np.uint8)            elif self.state == LoadState.loaded:                return self.data            elif self.state == LoadState.loadedCompressed:                raise Exception('Saliency maps do no have compressed handlind method enabled')                return None        elif self.lowImageType == InputType.fixationMapMatlab:            if self.state == LoadState.unloaded:                return (scipy.io.loadmat(self.filePath)['I']).astype(np.uint8)            elif self.state == LoadState.loaded:                return self.data            elif self.state == LoadState.loadedCompressed:                raise Exception('Fixation maps do no have compressed handlind method enabled')                return None        elif self.lowImageType == InputType.empty:            return None###############################################################################################class Target():    def __init__(self, lowImagePath, rawImagePath,                 lowImageState=LoadState.unloaded, lowImageType=InputType.image,                 rawImageState=LoadState.unloaded, rawImageType=InputType.saliencyMapMatlab):        self.lowImage = ImageContainer(lowImagePath, lowImageType, lowImageState)        self.rawImage = ImageContainer(rawImagePath, rawImageType, rawImageState)################################################################################################ class Target():#     def __init__(self, imagePath, saliencyPath,fixationPath,#                  imageState=LoadState.unloaded, imageType=InputType.image,#                  saliencyState=LoadState.unloaded, saliencyType=InputType.saliencyMapMatlab,#                  fixationState=LoadState.unloaded, fixationType=InputType.fixationMapMatlab):#         self.image = ImageContainer(imagePath, imageType, imageState)#         self.saliency = ImageContainer(saliencyPath, saliencyType, saliencyState)#         self.fixation = ImageContainer(fixationPath, fixationType, fixationState)# class Target1():#     def __init__(self, lowImagePath, rawImagePath,#                  lowImageState=LoadState.unloaded, lowImageType=InputType.image,#                  ):#         self.lowImage = ImageContainer(lowImagePath, lowImageType, lowImageState)#         self.rawImage = ImageContainer(rawImagePath, rawImageType, rawImageState)

第一个python代码调用了第二个python代码的。


希望对有需要的人有帮助。

0 0
原创粉丝点击