dataio.py 文件源码

def next(self):
        nrow=0
        ncol=0
        crop_size=self._crop_size
        while (nrow<crop_size or ncol<crop_size) \
              and self.cur_batch < self.batch_num:                  
            img_path=os.path.join(self._datadir, self._img_list[self.cur_batch])      
            img=cv2.imread(img_path, cv2.IMREAD_COLOR)
#            img=cv2.cvtColor(img,cv2.COLOR_BGR2YCR_CB)
            nrow,ncol=img.shape[0:2]
            self.cur_batch+=1
        if self.cur_batch < self.batch_num:
            sub_img_lr=npy.zeros(self._provide_data[0][1],dtype=npy.float32)
            sub_img_pryd=[]
            for item in self._provide_label:
                sub_img_pryd.append(npy.zeros(item[1],dtype=npy.float32))

            for i in range(self._crop_num):
                nrow_start=npy.random.randint(0,nrow-crop_size)
                ncol_start=npy.random.randint(0,ncol-crop_size)
                img_crop=img[nrow_start:nrow_start+crop_size,
                                ncol_start:ncol_start+crop_size,:] 
                imggt_size=crop_size                
                for s in range(self._num_scales):
                    img_temp=img_crop.astype(npy.float32)
                    img_temp=(img_temp-128)/128.0
                    img_temp = npy.swapaxes(img_temp, 0, 2)
                    img_temp = npy.swapaxes(img_temp, 1, 2)
                    sub_img_pryd[self._num_scales-s-1][i,:,:,:]=img_temp
                    imggt_size=imggt_size/2
                    img_crop=cv2.resize(img_crop,(imggt_size, imggt_size),
                              interpolation=cv2.INTER_CUBIC) 
                img_temp=img_crop.astype(npy.float32)
                img_temp=(img_temp-128)/128.0
                img_temp = npy.swapaxes(img_temp, 0, 2)
                img_temp = npy.swapaxes(img_temp, 1, 2)
                sub_img_lr[i,:,:,:]=img_temp              

            return LapSRNDataBatch(sub_img_lr,sub_img_pryd,0)
        else:
            raise StopIteration