'''
ResizeUtils提供了调整大小的功能，以保持长宽比完整。原作者：AndyP at StackOverflow'''
class ResizeUtils:
    # 给定目标高度，通过计算宽度和调整大小来调整图像
    def rescale_by_height(self, image, target_height,
                        method = cv2.INTER_LANCZOS4):
        # 将`image`重新缩放为`target_height`＃(保留宽高比)
        w = int(round(target_height * image.shape[1] / image.shape[0]))
        return (cv2.resize(image, (w, target_height),
                             interpolation = method))
    # 给定目标宽度，通过计算高度和调整大小来调整图像
    def rescale_by_width(self, image, target_width,
                        method = cv2.INTER_LANCZOS4):
        # 将`image`重新缩放为`target_width`(保留宽高比)
        h = int(round(target_width * image.shape[0] / image.shape[1]))
        return (cv2.resize(image, (target_width, h),
                            interpolation = method))

# FramesGenerator从给定的视频文件中提取图像帧
# 调整图像帧的大小以进行face_recognition / dlib处理
class FramesGenerator:
    def __init__(self, VideoFootageSource):
        self.VideoFootageSource = VideoFootageSource
    # 调整给定输入的大小以适合指定的尺寸以提取人脸嵌入
    def AutoResize(self, frame):
        resizeUtils = ResizeUtils()
        height, width, _ = frame.shape
        if height > 500:
            frame = resizeUtils.rescale_by_height(frame, 500)
            self.AutoResize(frame)
        if width > 700:
            frame = resizeUtils.rescale_by_width(frame, 700)
            self.AutoResize(frame)
        return frame

# 从视频素材＃中每秒提取1帧，并将这些帧保存到特定文件夹
def GenerateFrames(self, OutputDirectoryName):
    cap = cv2.VideoCapture(self.VideoFootageSource)
    _, frame = cap.read()
    fps = cap.get(cv2.CAP_PROP_FPS)
    TotalFrames = cap.get(cv2.CAP_PROP_FRAME_COUNT)
    print("[INFO] Total Frames ", TotalFrames, " @ ", fps, " fps")
    print("[INFO] Calculating number of frames per second")
    CurrentDirectory = os.path.curdir
    OutputDirectoryPath = os.path.join(
      CurrentDirectory, OutputDirectoryName)
    if os.path.exists(OutputDirectoryPath):
        shutil.rmtree(OutputDirectoryPath)
        time.sleep(0.5)
    os.mkdir(OutputDirectoryPath)
    CurrentFrame = 1
    fpsCounter = 0
    FrameWrittenCount = 1
    while CurrentFrame < TotalFrames:
        _, frame = cap.read()
        if (frame is None):
            continue
        if fpsCounter > fps:
            fpsCounter = 0
            frame = self.AutoResize(frame)
            filename = "frame_" + str(FrameWrittenCount) + ".jpg"
            cv2.imwrite(os.path.join(
              OutputDirectoryPath, filename), frame)
            FrameWrittenCount += 1
        fpsCounter += 1
        CurrentFrame += 1
    print('[INFO] Frames extracted')

# 以下是管道构造的节点。 
# 将创建并异步执行线程读取图像，提取面部特征并将其独立存储在不同的线程中
# 继续将文件名发送到管道中进行处理
class FramesProvider(Node):
    def setup(self, sourcePath):
        self.sourcePath = sourcePath
        self.filesList = []
        for item in os.listdir(self.sourcePath):
            _, fileExt = os.path.splitext(item)
            if fileExt == '.jpg':
                self.filesList.append(os.path.join(item))
        self.TotalFilesCount = self.size = len(self.filesList)
        self.ProcessedFilesCount = self.pos = 0
    # 发出管道中的每个文件名以进行并行处理
    def run(self, data):
        if self.ProcessedFilesCount < self.TotalFilesCount:
            self.emit({'id': self.ProcessedFilesCount,
                'imagePath': os.path.join(self.sourcePath,
                              self.filesList[self.ProcessedFilesCount])})
            self.ProcessedFilesCount += 1
            self.pos = self.ProcessedFilesCount
        else:
            self.close()

# 对面部嵌入，参考路径号和位置进行编码，然后发送到管道
class FaceEncoder(Node):
    def setup(self, detection_method = 'cnn'):
        self.detection_method = detection_method
        # detection_method可以是cnn或hog
    def run(self, data):
        id = data['id']
        imagePath = data['imagePath']
        image = cv2.imread(imagePath)
        rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        boxes = face_recognition.face_locations(
               rgb, model = self.detection_method)
        encodings = face_recognition.face_encodings(rgb, boxes)
        d = [{"imagePath": imagePath, "loc": box, "encoding": enc}
                         for (box, enc) in zip(boxes, encodings)]
        self.emit({'id': id, 'encodings': d})

# 接收面部encoding以进行聚类，并获id来命名不同的文件名
class DatastoreManager(Node):
    def setup(self, encodingsOutputPath):
        self.encodingsOutputPath = encodingsOutputPath
    def run(self, data):
        encodings = data['encodings']
        id = data['id']
        with open(os.path.join(self.encodingsOutputPath,
                   'encodings_' + str(id) + '.pickle'), 'wb') as f:
            f.write(pickle.dumps(encodings))

# PicklesListCollat​​or将多个pickle文件作为输入并将它们合并在一起。
# 它专门用于支持将不同的pickle文件合并为一个的用例
class PicklesListCollator:
    def __init__(self, picklesInputDirectory):
        self.picklesInputDirectory = picklesInputDirectory
    # 在这里，我们将列出所有pickle
    # 从多个线程生成的文件，读取结果列表，将它们追加到公用列表中，并创建另一个以组合列表为内容的pickle
    def GeneratePickle(self, outputFilepath):
        datastore = []
        ListOfPickleFiles = []
        for item in os.listdir(self.picklesInputDirectory):
            _, fileExt = os.path.splitext(item)
            if fileExt == '.pickle':
                ListOfPickleFiles.append(os.path.join(
                    self.picklesInputDirectory, item))
        for picklePath in ListOfPickleFiles:
            with open(picklePath, "rb") as f:
                data = pickle.loads(f.read())
                datastore.extend(data)
        with open(outputFilepath, 'wb') as f:
            f.write(pickle.dumps(datastore))

# 人脸聚类功能
class FaceClusterUtility:
    def __init__(self, EncodingFilePath):
        self.EncodingFilePath = EncodingFilePath
    # 鸣谢：Arian的pyimagesearch用于聚类代码
    # 在这里，我们使用sklearn.DBSCAN函数
    # 聚类所有面部encoding以获得代表不同的人的聚类
    def Cluster(self):
        InputEncodingFile = self.EncodingFilePath
        if not (os.path.isfile(InputEncodingFile) and
                os.access(InputEncodingFile, os.R_OK)):
            print('The input encoding file, ' +
                    str(InputEncodingFile) +
                    ' does not exists or unreadable')
            exit()
        NumberOfParallelJobs = -1
        # 从磁盘加载序列化的面部编码+边界框位置，然后将编码集提取，以便我们可以在它们上聚类
        print("[INFO] Loading encodings")
        data = pickle.loads(open(InputEncodingFile, "rb").read())
        data = np.array(data)
        encodings = [d["encoding"] for d in data]
        # 将embeddings聚类
        print("[INFO] Clustering")
        clt = DBSCAN(eps = 0.5, metric ="euclidean",
                      n_jobs = NumberOfParallelJobs)
        clt.fit(encodings)
        # 确定在数据集中发现的独特面孔的总数
        labelIDs = np.unique(clt.labels_)
        numUniqueFaces = len(np.where(labelIDs > -1)[0])
        print("[INFO] # unique faces: {}".format(numUniqueFaces))
        return clt.labels_

# 继承类tqdm以可视化进度
class TqdmUpdate(tqdm):
    # 此函数将作为progress回调函数传递。为可视化中的自动更新设置预定义的变量
    def update(self, done, total_size = None):
        if total_size is not None:
            self.total = total_size
        self.n = done
        super().refresh()

class FaceImageGenerator:
    def __init__(self, EncodingFilePath):
        self.EncodingFilePath = EncodingFilePath
    # 在这里，我们为每个不同的面孔创建前25个面孔的蒙太奇。 
    # 我们还将使用簇中的标签和编码和pickle文件中的图像url为所有个不同的面孔生成图像。 
    # 为训练目的增加了面部边界框，我们还为每个面部图像创建了精确的注释(类似于Darknet YOLO)
    # 为了将来的使用轻松地修改注释。并能够在监督训练中使用
    def GenerateImages(self, labels, OutputFolderName = "ClusteredFaces",
                                            MontageOutputFolder = "Montage"):
        output_directory = os.getcwd()
        OutputFolder = os.path.join(output_directory, OutputFolderName)
        if not os.path.exists(OutputFolder):
            os.makedirs(OutputFolder)
        else:
            shutil.rmtree(OutputFolder)
            time.sleep(0.5)
            os.makedirs(OutputFolder)
        MontageFolderPath = os.path.join(OutputFolder, MontageOutputFolder)
        os.makedirs(MontageFolderPath)
        data = pickle.loads(open(self.EncodingFilePath, "rb").read())
        data = np.array(data)
        labelIDs = np.unique(labels)
        # 遍历唯一的面孔整数
        for labelID in labelIDs:
            # 在属于当前标签ID的data数组中查找所有索引，然后从集合中随机抽取最多25个索引
            print("[INFO] faces for face ID: {}".format(labelID))
            FaceFolder = os.path.join(OutputFolder, "Face_" + str(labelID))
            os.makedirs(FaceFolder)
            idxs = np.where(labels == labelID)[0]
            # 将面孔列表初始化为，包含在蒙太奇中
            portraits = []
            # 循环采样索引
            counter = 1
            for i in idxs:
                # 加载输入图像并提取面部ROI
                image = cv2.imread(data[i]["imagePath"])
                (o_top, o_right, o_bottom, o_left) = data[i]["loc"]
                height, width, channel = image.shape
                widthMargin = 100
                heightMargin = 150
                top = o_top - heightMargin
                if top < 0: top = 0
                bottom = o_bottom + heightMargin
                if bottom > height: bottom = height
                left = o_left - widthMargin
                if left < 0: left = 0
                right = o_right + widthMargin
                if right > width: right = width
                portrait = image[top:bottom, left:right]
                if len(portraits) < 25:
                    portraits.append(portrait)
                resizeUtils = ResizeUtils()
                portrait = resizeUtils.rescale_by_width(portrait, 400)
                FaceFilename = "face_" + str(counter) + ".jpg"
                FaceImagePath = os.path.join(FaceFolder, FaceFilename)
                cv2.imwrite(FaceImagePath, portrait)
                widthMargin = 20
                heightMargin = 20
                top = o_top - heightMargin
                if top < 0: top = 0
                bottom = o_bottom + heightMargin
                if bottom > height: bottom = height
                left = o_left - widthMargin
                if left < 0: left = 0
                right = o_right + widthMargin
                if right > width:
                    right = width
                AnnotationFilename = "face_" + str(counter) + ".txt"
                AnnotationFilePath = os.path.join(FaceFolder, AnnotationFilename)
                f = open(AnnotationFilePath, 'w')
                f.write(str(labelID) + ' ' +
                        str(left) + ' ' + str(top) + ' ' +
                        str(right) + ' ' + str(bottom) + "\n")
                f.close()
                counter += 1
            montage = build_montages(portraits, (96, 120), (5, 5))[0]
            MontageFilenamePath = os.path.join(
               MontageFolderPath, "Face_" + str(labelID) + ".jpg")
            cv2.imwrite(MontageFilenamePath, montage)

# 从上面的Python文件导入所有类
from FaceClusteringLibrary import *
if __name__ == "__main__":
    # 从给定的视频素材生成帧
    framesGenerator = FramesGenerator("Footage.mp4")
    framesGenerator.GenerateFrames("Frames")
    # 设计并运行人脸聚类管道
    CurrentPath = os.getcwd()
    FramesDirectory = "Frames"
    FramesDirectoryPath = os.path.join(CurrentPath, FramesDirectory)
    EncodingsFolder = "Encodings"
    EncodingsFolderPath = os.path.join(CurrentPath, EncodingsFolder)
    if os.path.exists(EncodingsFolderPath):
        shutil.rmtree(EncodingsFolderPath, ignore_errors = True)
        time.sleep(0.5)
    os.makedirs(EncodingsFolderPath)
    pipeline = Pipeline(
                    FramesProvider("Files source", sourcePath = FramesDirectoryPath) |
                    FaceEncoder("Encode faces") |
                    DatastoreManager("Store encoding",
                    encodingsOutputPath = EncodingsFolderPath),
                    n_threads = 3, quiet = True)
    pbar = TqdmUpdate()
    pipeline.run(update_callback = pbar.update)
    print()
    print('[INFO] Encodings extracted')
    # 将所有编码pickle文件合并为一个
    CurrentPath = os.getcwd()
    EncodingsInputDirectory = "Encodings"
    EncodingsInputDirectoryPath = os.path.join(
          CurrentPath, EncodingsInputDirectory)
    OutputEncodingPickleFilename = "encodings.pickle"
    if os.path.exists(OutputEncodingPickleFilename):
        os.remove(OutputEncodingPickleFilename)
    picklesListCollator = PicklesListCollator(
                    EncodingsInputDirectoryPath)
    picklesListCollator.GeneratePickle(
           OutputEncodingPickleFilename)
    # 管理文件写入中的任何延迟
    time.sleep(0.5)
    # 开始聚类过程并生成带有注释的输出图像
    EncodingPickleFilePath = "encodings.pickle"
    faceClusterUtility = FaceClusterUtility(EncodingPickleFilePath)
    faceImageGenerator = FaceImageGenerator(EncodingPickleFilePath)
    labelIDs = faceClusterUtility.Cluster()
    faceImageGenerator.GenerateImages(
      labelIDs, "ClusteredFaces", "Montage")