#此模块主要用来读取数据集，返回一个数据可迭代对象
#重点是，此模块分批次的把图像读入内存的，而不是一次全读入，有效的减少了内存溢出
import os
import cv2
import numpy as np
import keras
from random import shuffle

#目标图像大小
image_size= (229, 229, 3)
#类别编码
class_dict=dict(zip(['neg','pos','neu'],[0,1,2]))
#处理.txt文件,并加载图片文件夹里的图片名
#txt_path，txt文件路径，data_path，图片文件夹路径

def read_txt(txt_path,data_path):
 # 中间数组
 labels_class = []
 labels_score = []
 with open(txt_path) as f:
 lines_list = f.readlines()
 for lines in lines_list:
  line = lines.split(' ')
  labels_class.append(line[0].rstrip(".jpg"))
  labels_score.append(line[2])
 labels_dict=dict(zip(labels_class,labels_score))
 #处理图片数据集
 all_picture_name = os.listdir(data_path)
 #打乱数据集
 shuffle(all_picture_name)
 all_picture_path=[os.path.join(data_path,one)for one in all_picture_name]
 return all_picture_name,all_picture_path,labels_dict

class data_generate:
 def __init__(self,all_piture_name,all_picture_path,labels_dict,batch_size):
 self.index=0
 self.all_picture_name=all_piture_name
 self.all_picture_path=all_picture_path
 self.labels_dict=labels_dict
 self.batch_size = batch_size
 def get_mini_batch(self):
  while True:
  batch_images=[]
  batch_labels=[]
  batch_class=[]
  batch_score=[]
  for i in range(self.batch_size):
  if(self.index==len(self.all_picture_name)):
   self.index=0

  bgr_image = cv2.imread(self.all_picture_path[self.index])
  if len(bgr_image.shape) == 2: # 若是灰度图则转为三通道
   bgr_image = cv2.cvtColor(bgr_image, cv2.COLOR_GRAY2BGR)
  rgb_image = cv2.cvtColor(bgr_image, cv2.COLOR_BGR2RGB)
  rgb_image=cv2.resize(rgb_image,(image_size[0], image_size[1]))
  img = np.array(rgb_image)
  img=keras.applications.inception_v3.preprocess_input(img)
  batch_images.append(img)
  #label=[]
  label1=self.all_picture_name[self.index].rstrip(".jpg")
  batch_class.append(keras.utils.to_categorical(class_dict[label1[:3]], 3))
  batch_score.append(np.array(self.labels_dict[label1]))
  #batch_labels.append(label)
  self.index+=1
  batch_images=np.array(batch_images)
  batch_class = np.array(batch_class)
  batch_score = np.array(batch_score)
  #注意label的生成batch_class，一个单独数组，batch_score一个单独的数组
  '''
  注释掉的这段代码生成的label是错误的。
  batch_images=[]
  batch_labels=[]
  for i in range(self.batch_size):
  if(self.index==len(self.images)):
   self.index=0
  batch_images.append(self.images[self.index])
  batch_labels.append(self.labels[self.index])
  self.index+=1
  batch_images=np.array(batch_images)
  batch_labels=np.array(batch_labels)
  yield batch_images,batch_labels
  '''
  yield batch_images,[batch_class,batch_score]

history=model.fit_generator(generator=train_data.get_mini_batch(),
   steps_per_epoch=146,
   epochs=300,
   validation_data=test_data.get_mini_batch(),
   validation_steps=34,
   )