import numpy as np
import cv2
import os
from utility import Method
from image_fusion import ImageFusion
import time
import copy
from images_stitcher import ImagesStitch

class VideoStitch(Method):
    """
    The class of video stitcher
    """
    image_shape = None
    offset_list = []
    is_available_list = []
    images_address_list = None

    need_detailed_register = True

    # 关于融合方法的设置
    fuse_method = "not_fuse"
    sample_rate = 1
    video_name = ""
    image_stitcher = ImagesStitch()

    def start_stitching(self, video_address, sample_rate=1, use_pre_calculate=False,
                        pre_calculate_available = None, pre_calculate_offset=None, pre_register_time=None):
        """
        stitching the video
        :param video_address: 视频地址
        :param sample_rate: 视频采样帧率
        :param use_pre_calculate: 是否使用预先计算的偏移量
        :param pre_calculate_available: 预先计算的可用list
        :param pre_calculate_offset: 预先计算的的偏移量
        :param pre_register_times: 预先计算的配准时间
        :return: 返回拼接后的图像，ndarry
        """
        # *********** 对视频采样，将采样的所有图像输出到与视频文件同目录的temp文件夹 ***********
        # 建立 temp 文件夹
        input_dir = os.path.dirname(video_address)
        video_name = os.path.basename(video_address).split(".")[0]
        self.video_name = video_name
        sample_dir = os.path.join(input_dir, "temp_" + video_name)

        # 将 video 采样到 temp 文件夹
        self.print_and_log("Video name:" + video_address)
        self.print_and_log("Sampling rate:" + str(sample_rate))
        self.print_and_log("We save sampling images in " + sample_dir)
        self.print_and_log("Sampling images ...")

        if os.path.exists(sample_dir):
            self.delete_folder(sample_dir)
        self.make_out_dir(sample_dir)

        # 解压文件时有可能会得到无法解压的错误，需要在工程中注意
        self.sample_rate = sample_rate
        cap = cv2.VideoCapture(video_address)
        frame_num = 0
        save_num = 0
        start_time = time.time()
        while True:
            ret, origin_frame = cap.read()
            if ret is False:
                break
            frame_num = frame_num + 1
            if frame_num % sample_rate == 0:
                # 由于视频采集最后一行默认为0，所以需要去掉
                gray_frame = cv2.cvtColor(origin_frame, cv2.COLOR_BGR2GRAY)[:-2, :]
                save_num = save_num + 1
                cv2.imwrite(os.path.join(sample_dir, str(save_num).zfill(10) + ".png"), gray_frame)
        cap.release()
        end_time = time.time()
        self.print_and_log("Sampled done, The time of sampling is {:.3f} \'s".format(end_time - start_time))

        # # **************************** 配准 ****************************
        # # 开始拼接文件夹下的图片
        dirs = sorted(os.listdir(sample_dir), key=lambda i: int(i.split(".")[0]))
        self.images_address_list = [os.path.join(sample_dir, item) for item in dirs]
        temp_image = cv2.imdecode(np.fromfile(self.images_address_list[0], dtype=np.uint8),
                                  cv2.IMREAD_GRAYSCALE)
        self.image_shape = temp_image.shape
        self.print_and_log("start matching")
        del temp_image
        start_time = time.time()
        if use_pre_calculate:
            self.is_available_list = pre_calculate_available
            self.offset_list = pre_calculate_offset
        else:
            self.is_available_list.append(True)
            for file_index in range(1, len(self.images_address_list)):
                self.print_and_log("    Analyzing {}th frame and the name is {}".format(file_index, os.path.basename(
                    self.images_address_list[file_index])))
                # 获得上一次可用的编号,不是所有帧都是有用的，所以在寻找时需要过滤
                # 先找到 is_available_list 中最后一个True的索引，在减去目前索引到其距离
                temp_available_list = self.is_available_list.copy()
                temp_available_list.reverse()
                last_file_index = file_index - temp_available_list.index(True) - 1
                # self.print_and_log("  The last file index is {}, the next file index is {}"
                #                    .format(last_file_index, file_index))
                last_image = cv2.imdecode(np.fromfile(self.images_address_list[last_file_index], dtype=np.uint8),
                                          cv2.IMREAD_GRAYSCALE)
                next_image = cv2.imdecode(np.fromfile(self.images_address_list[file_index], dtype=np.uint8),
                                          cv2.IMREAD_GRAYSCALE)
                self.image_shape = last_image.shape
                status, offset = self.calculate_offset_by_feature_in_roi([last_image, next_image])

                if status and self.need_detailed_register:
                    overlap_last_image, overlap_next_image = \
                        self.get_overlap_regions(last_image, next_image, offset)
                    # print(overlap_last_image.shape)
                    # print(overlap_next_image.shape)
                    detailed_status, detailed_offset = \
                        self.calculate_offset_by_feature_in_roi([overlap_last_image, overlap_next_image])
                    if detailed_status:
                        offset[0] = offset[0] + detailed_offset[0]
                        offset[1] = offset[1] + detailed_offset[1]
                if status is False:
                    self.print_and_log("    {}th frame can not be stitched, the reason is {}".format(file_index, offset))
                    self.is_available_list.append(False)
                    self.offset_list.append([0, 0])
                else:
                    self.print_and_log("    {}th frame can be stitched, the offset is {}".format(file_index, offset))
                    self.is_available_list.append(True)
                    self.offset_list.append(offset)
                del last_image, next_image
        record_offset_list = copy.deepcopy(self.offset_list)

        end_time = time.time()
        if use_pre_calculate:
            self.print_and_log("The time of registering is {:.3f} \'s".format(end_time - start_time + pre_register_time))
        else:
            self.print_and_log("The time of registering is {:.3f} \'s".format(end_time - start_time))
        self.print_and_log("is_available_list:{}".format(self.is_available_list))
        self.print_and_log("offset_list:{}".format(self.offset_list))
        self.print_and_log("len is_available_list:{}".format(len(self.is_available_list)))
        self.print_and_log("len offset_list:{}".format(len(self.offset_list)))
        # self.offset_list = [[1, 0], [-2, 0], [1, 0], [-1, 0], [1, 0], [0, -1], [3, 1], [0, -1], [0, 1], [0, -1], [0, 1], [0, -1], [-1, 0], [-1, 0], [1, 0], [1, 0], [1, 0], [8, 2], [11, 1], [33, 6], [34, 7], [69, 14], [12, 2], [2, 0], [2, 0], [0, -1], [-1, 0], [0, 1], [-1, 0], [-1, 0], [-1, 0], [-1, 0], [1, 0], [1, 0], [1, 0], [1, 0], [1, 0], [5, 1], [2, 0], [3, 0], [12, 2], [0, 0], [87, 18], [7, 1], [3, 0], [2, 0], [1, 0], [1, 0], [0, -1], [0, -1], [0, 1], [-1, 0], [0, -1], [-1, 0], [0, 1], [1, 0], [0, 1], [0, 1], [0, -1], [1, 0], [1, 0], [1, 0], [1, 0], [1, 0], [0, 1], [0, 1], [1, 0], [1, 0], [1, 0], [1, 0], [1, 0], [4, 1], [4, 1], [2, 0], [1, 0], [1, 0], [5, -3], [4, -3], [1, 0], [1, 0], [-3, 0], [1, 0], [4, 1], [2, 0], [4, 0], [2, 0], [4, 0], [3, 1], [5, 0], [7, -2], [6, -1], [-7, -1], [1, 0], [2, 0], [0, 1], [-1, -1], [-1, 0], [0, 1], [0, -1], [1, 0], [-1, 0], [0, -1], [1, 0], [0, 1], [1, 0], [1, 0], [10, 2], [4, 0], [1, 0], [1, 0], [0, -1], [3, 1], [1, 0], [-4, -1], [-1, 0], [1, 0], [1, 0], [2, 0], [4, 1], [6, 1], [27, 0], [0, 0], [258, 39], [23, -13], [0, 1], [0, -1], [1, 0], [-1, 0], [-1, 0], [-1, 0], [-1, 0], [1, 0], [2, 0], [1, 0], [0, -1], [0, 1], [-1, 0], [1, 0], [1, 0], [1, 0], [0, -1], [-1, 0], [1, 0], [-1, 0], [1, 0], [4, 0], [0, -1], [0, -1], [1, 0], [12, -11], [1, 0], [-1, 0], [-1, 0], [-2, 0], [-1, 0], [-3, 0], [-2, 0], [0, 1], [0, -1], [0, 1], [-1, 0], [0, -1], [0, 1], [0, 1], [1, 0], [-1, 0], [-1, 0], [0, -1], [0, 1], [1, 0], [1, 0], [0, 1], [-1, 0], [0, 1], [1, 0], [0, -1], [1, 0], [0, 1], [0, -1], [1, 0], [0, 1], [1, 0], [8, 1], [1, 0], [3, -1], [4, -2], [9, -7], [4, -6], [0, -2], [-7, -3], [7, 1], [2, 0], [-5, -1], [1, 0], [1, 0], [1, 0], [-1, 0], [1, 0], [-1, 0], [-6, -1], [2, 0], [12, -3], [3, -2], [6, -6], [3, -2], [6, -5], [10, -9], [6, -5], [2, -3], [1, 0], [2, 0], [4, -3], [17, -15], [1, -1], [1, 0], [2, 0], [2, -3], [2, 0], [1, 0], [0, -1], [1, 0], [1, 0], [1, 0], [-2, 0], [-7, -1], [1, 0], [1, 0], [1, 0], [0, -1], [1, 0], [4, 1], [9, 1], [6, 1], [9, 1], [17, 4], [20, 4], [11, 2], [7, 1], [4, 1], [0, -1], [1, 0], [0, -1], [0, -1], [1, 0], [0, 1], [0, 1], [1, 0], [1, 0], [1, 0], [0, 1], [1, 0], [0, -1], [0, 1], [0, -1], [1, 0], [1, 0], [0, -1], [1, 0], [-1, 0], [0, 1], [1, 0], [1, 0], [1, 0], [1, 0], [-1, 0], [1, 0], [1, 0], [0, 1], [1, 0], [1, 0], [0, 1], [0, -1], [-1, 0], [0, 1], [1, 0], [1, 0], [1, 0], [0, 1], [0, -1], [0, -1], [0, 1], [-1, 0], [1, 0], [0, 1], [0, -1], [0, -1], [1, 0], [0, 1], [0, 1], [0, -1], [1, 0], [1, 1], [3, 0], [0, 0], [0, 0], [258, 57], [14, 3], [7, 0], [1, 0], [-6, -1], [-7, -1], [-1, 0], [3, 0], [5, 0], [-3, 0], [-1, 0], [6, 0], [-1, 0], [-7, -1], [-1, 0], [-1, 0], [1, 0], [0, -1], [0, 1], [0, 1], [0, 1], [0, 1], [0, -1], [0, -1], [0, -1], [0, -1], [0, 1], [0, 1], [0, -1], [-1, 0], [0, 1], [0, -1], [0, 1], [0, -1], [0, -1], [0, 1], [0, 1], [1, 0], [1, 0], [0, 1], [0, 1], [-1, 0], [0, 1], [0, -1], [0, -1], [0, 1], [1, 0], [0, -1], [0, 1], [0, -1], [0, 1], [-1, 0], [0, -1], [8, 1], [5, 1], [0, 1], [0, -1], [0, 1], [0, 1], [0, -1], [0, 1], [0, 1], [-1, 0], [0, -1], [0, 1], [0, -1], [1, 0], [1, 0], [0, -1], [-1, 0], [0, -1], [0, -1], [-1, 0], [1, 0], [1, 0], [0, 1], [0, -1], [0, -1], [0, 1]]
        # self.is_available_list = [True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, False, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, False, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, False, False, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True]
        # # *************************** 融合及拼接 ***************************
        self.print_and_log("start fusing")
        start_time = time.time()
        stitch_image = self.get_stitch_by_offset()
        end_time = time.time()
        self.print_and_log("The time of fusing is {:.3f} \'s".format(end_time - start_time))
        self.delete_folder(sample_dir)
        record_available_list = self.is_available_list
        self.is_available_list = []
        self.offset_list = []
        return stitch_image, record_available_list, record_offset_list

    def calculate_offset_by_feature_in_roi(self, images):
        """
        采用ROI增长的方式进行特征搜索
        :param images: [last_image, next_image]
        :return: status, temp_offset, （拼接状态， 偏移量）
        """
        status = False
        temp_offset = [0, 0]
        last_image, next_image = images
        row, col = last_image.shape[:2]
        max_iteration = int(1 / self.roi_ratio)
        for i in range(0, max_iteration):
            temp_ratio = self.roi_ratio * (i + 1)
            search_len = np.floor(row * temp_ratio).astype(int)
            roi_last_image = last_image[0: search_len, :]
            roi_next_image = next_image[0: search_len, :]
            status, temp_offset = self.calculate_offset_by_feature(roi_last_image, roi_next_image)
            if status is False:
                continue
            else:
                break
        return status, temp_offset

    def calculate_offset_by_feature(self, last_image, next_image):
        """
        通过全局特征匹配计算偏移量
        :param last_image: 上一张图像
        :param next_image: 下一张图像
        :return: status, temp_offset, （拼接状态， 偏移量）
        """
        offset = [0, 0]
        status = False

        # get the feature points
        last_kps, last_features = self.calculate_feature(last_image)
        next_kps, next_features = self.calculate_feature(next_image)
        if last_features is not None and next_features is not None:
            matches = self.match_descriptors(last_features, next_features)
            # match all the feature points
            if self.offset_calculate == "mode":
                (status, offset) = self.get_offset_by_mode(last_kps, next_kps, matches)
            elif self.offset_calculate == "ransac":
                (status, offset, adjustH) = self.get_offset_by_ransac(last_kps, next_kps, matches)
            return status, offset
        else:
            return status, "there are one image have no features"

    def get_stitch_by_offset(self):
        """
        根据偏移量计算返回拼接结果
        :return: 拼接结果图像
        """
        # 如果你不细心，不要碰这段代码
        # 已优化到根据指针来控制拼接，CPU下最快了
        dx_sum = dy_sum = 0
        result_row = self.image_shape[0]  # 拼接最终结果的横轴长度,先赋值第一个图像的横轴
        result_col = self.image_shape[1]  # 拼接最终结果的纵轴长度,先赋值第一个图像的纵轴
        self.offset_list.insert(0, [0, 0])  # 增加第一张图像相对于最终结果的原点的偏移量
        range_x = [[0, 0] for _ in range(len(self.offset_list))]  # 主要用于记录X方向最大最小边界
        range_y = [[0, 0] for _ in range(len(self.offset_list))]  # 主要用于记录Y方向最大最小边界
        temp_offset_list = self.offset_list.copy()
        offset_list_origin = self.offset_list.copy()
        range_x[0][1] = self.image_shape[0]
        range_y[0][1] = self.image_shape[1]

        for i in range(1, len(temp_offset_list)):
            if self.is_available_list[i] is False:
                range_x[i][0] = range_x[i - 1][0]
                range_x[i][1] = range_x[i - 1][1]
                range_y[i][0] = range_y[i - 1][0]
                range_y[i][1] = range_y[i - 1][1]
                continue
            dx_sum = dx_sum + self.offset_list[i][0]
            dy_sum = dy_sum + self.offset_list[i][1]
            if dx_sum <= 0:
                for j in range(0, i):
                    temp_offset_list[j][0] = temp_offset_list[j][0] + abs(dx_sum)
                    range_x[j][0] = range_x[j][0] + abs(dx_sum)
                    range_x[j][1] = range_x[j][1] + abs(dx_sum)
                result_row = result_row + abs(dx_sum)
                range_x[i][1] = result_row
                dx_sum = range_x[i][0] = temp_offset_list[i][0] = 0
            else:
                temp_offset_list[i][0] = dx_sum
                result_row = max(result_row, dx_sum + self.image_shape[0])
                range_x[i][1] = result_row
            if dy_sum <= 0:
                for j in range(0, i):
                    temp_offset_list[j][1] = temp_offset_list[j][1] + abs(dy_sum)
                    range_y[j][0] = range_y[j][0] + abs(dy_sum)
                    range_y[j][1] = range_y[j][1] + abs(dy_sum)
                result_col = result_col + abs(dy_sum)
                range_y[i][1] = result_col
                dy_sum = range_y[i][0] = temp_offset_list[i][1] = 0
            else:
                temp_offset_list[i][1] = dy_sum
                result_col = max(result_col, dy_sum + self.image_shape[1])
                range_y[i][1] = result_col
        stitch_result = np.zeros((result_row, result_col), np.int) - 1
        self.offset_list = temp_offset_list
        self.print_and_log("  The rectified offsetList is " + str(self.offset_list))
        # print(len(self.offset_list))
        # 如上算出各个图像相对于原点偏移量，并最终计算出输出图像大小，并构造矩阵，如下开始赋值
        count = 0
        for i in range(0, len(self.offset_list)):
            if self.is_available_list[i] is False:
                continue
            image = cv2.imdecode(np.fromfile(self.images_address_list[i], dtype=np.uint8), cv2.IMREAD_GRAYSCALE)
            self.print_and_log("  stitching " + str(self.images_address_list[i]))
            if i == 0:
                stitch_result[self.offset_list[0][0]: self.offset_list[0][0] + image.shape[0],
                              self.offset_list[0][1]: self.offset_list[0][1] + image.shape[1]] = image
            else:
                count += 1
                min_occupy_x = range_x[i - 1][0]
                max_occupy_x = range_x[i - 1][1]
                min_occupy_y = range_y[i - 1][0]
                max_occupy_y = range_y[i - 1][1]
                # self.printAndWrite("Stitch " + str(i + 1) + "th, the offsetList[i][0] is " + str(
                #     offsetList[i][0]) + " and the offsetList[i][1] is " + str(offsetList[i][1]))
                # self.printAndWrite("Stitch " + str(i + 1) + "th, the minOccupyX is " + str(
                #     minOccupyX) + " and the maxOccupyX is " + str(maxOccupyX) + " and the minOccupyY is " + str(
                #     minOccupyY) + " and the maxOccupyY is " + str(maxOccupyY))
                roi_ltx = max(self.offset_list[i][0], min_occupy_x)
                roi_lty = max(self.offset_list[i][1], min_occupy_y)
                roi_rbx = min(self.offset_list[i][0] + image.shape[0], max_occupy_x)
                roi_rby = min(self.offset_list[i][1] + image.shape[1], max_occupy_y)
                # self.printAndWrite("Stitch " + str(i + 1) + "th, the roi_ltx is " + str(
                #     roi_ltx) + " and the roi_lty is " + str(roi_lty) + " and the roi_rbx is " + str(
                #     roi_rbx) + " and the roi_rby is " + str(roi_rby))
                last_roi_fuse_region = stitch_result[roi_ltx:roi_rbx, roi_lty:roi_rby].copy()
                stitch_result[
                    self.offset_list[i][0]: self.offset_list[i][0] + image.shape[0],
                    self.offset_list[i][1]: self.offset_list[i][1] + image.shape[1]] = image.copy()
                next_roi_fuse_region = stitch_result[roi_ltx:roi_rbx, roi_lty:roi_rby].copy()
                stitch_result[roi_ltx:roi_rbx, roi_lty:roi_rby] = self.fuse_image(
                    [last_roi_fuse_region, next_roi_fuse_region],
                    [offset_list_origin[i][0], offset_list_origin[i][1]])
                self.save_roi_images(count, last_roi_fuse_region, next_roi_fuse_region)
                del last_roi_fuse_region, next_roi_fuse_region
            del image
        stitch_result[stitch_result == -1] = 0
        return stitch_result.astype(np.uint8)

    def fuse_image(self, overlap_rfrs, offset):
        """
        融合两个重合区域,其中rfr代表（roi_fuse_region）
        :param overlap_rfrs:重合区域
        :param offset: 原本两图像的位移
        :return:返回融合结果
        """
        (last_rfr, next_rfr) = overlap_rfrs
        if self.fuse_method != "fade_in_fade_out" and self.fuse_method != "trigonometric":
            # 将各自区域中为背景的部分用另一区域填充，目的是消除背景
            # 权值为-1是为了方便渐入检出融合和三角融合计算
            last_rfr[last_rfr == -1] = 0
            next_rfr[next_rfr == -1] = 0
            last_rfr[last_rfr == 0] = next_rfr[last_rfr == 0]
            next_rfr[next_rfr == 0] = last_rfr[next_rfr == 0]
        fuse_region = np.zeros(last_rfr.shape, np.uint8)
        image_fusion = ImageFusion()
        if self.fuse_method == "not_fuse":
            fuse_region = next_rfr
        elif self.fuse_method == "average":
            fuse_region = image_fusion.fuse_by_average([last_rfr, next_rfr])
        elif self.fuse_method == "maximum":
            fuse_region = image_fusion.fuse_by_maximum([last_rfr, next_rfr])
        elif self.fuse_method == "minimum":
            fuse_region = image_fusion.fuse_by_minimum([last_rfr, next_rfr])
        elif self.fuse_method == "fade_in_fade_out":
            fuse_region = image_fusion.fuse_by_fade_in_and_fade_out(overlap_rfrs, offset)
        elif self.fuse_method == "trigonometric":
            fuse_region = image_fusion.fuse_by_trigonometric(overlap_rfrs, offset)
        elif self.fuse_method == "spatial_frequency":
            fuse_region = image_fusion.fuse_by_spatial_frequency([last_rfr, next_rfr])
        elif self.fuse_method == "multi_band_blending":
            fuse_region = image_fusion.fuse_by_multi_band_blending([last_rfr, next_rfr])
        elif self.fuse_method == "our_framework":
            fuse_region = image_fusion.fuse_by_our_framework([last_rfr, next_rfr])
        return fuse_region

    @staticmethod
    def get_overlap_regions(last_image, next_image, offset):
        h, w = last_image.shape
        overlap_last_image = None
        overlap_next_image = None
        dx, dy = offset
        # print(offset)
        if dx >= 0 and dy >= 0:
            overlap_last_image = last_image[dx: h, dy: w]
            overlap_next_image = next_image[0: h - dx, 0: w - dy]
        elif dx >= 0 and dy < 0:
            overlap_last_image = last_image[dx: h, 0: w + dy]
            overlap_next_image = next_image[0: h - dx, -dy: w]
        elif dx < 0 and dy >= 0:
            overlap_last_image = last_image[0: h + dx, dy: w]
            overlap_next_image = next_image[-dx: h, 0: w - dy]
        elif dx < 0 and dy < 0:
            overlap_last_image = last_image[0: h + dx, 0: w + dy]
            overlap_next_image = next_image[-dx: h, -dy: w]
        return overlap_last_image, overlap_next_image

    def record_video_stitch_parameters(self, output_dir, video_name, available_list, offset_list, register_time):
        '''
        Write parameters in video stitch
        :param output_dir: the address of record file
        :param video_name: str
        :param available_list: like [True, False]
        :param offset_list: like [[20, 30], [-1, -2]]
        :param register_time: float, time of finding feature and matching
        :return:
        '''
        file_address = os.path.join(output_dir, "video_stitch_sample_" + str(self.sample_rate) + "_record.txt")
        f = open(file_address, "a")
        f.write("###-" + video_name)
        f.write("\n")
        temp = [str(1) if item else str(0) for item in available_list]
        f.write(",".join(temp))
        f.write("\n")
        temp = ""
        for item in offset_list:
            temp = temp + str(item[0]) + " " + str(item[1]) + ","
        f.write(temp[:-1])
        f.write("\n")
        f.write(str(register_time))
        f.write("\n")
        f.close()

    def read_video_stitch_parameters(self, input_address):
        '''
        Read the record with video stitch parameters
        :param input_address: record file address
        :return: available_list, offset_list, register_time_list
        '''
        available_list = []
        offset_list = []
        register_time_list =[]
        if os.path.exists(input_address) is False:
            self.print_and_log("There is no record file with video parameters to read")
            return available_list, offset_list, register_time_list
        count = 0   # using count to judge which value to read, 0 is available_list, 1 is offset_list and 2 is time
        with open(input_address, 'r') as f:
            while True:
                lines = f.readline()
                if not lines:
                    break
                if lines.startswith("###-"):
                    count = 0
                    continue
                lines = lines.strip('\n')
                if count == 0:
                    temp_available = []
                    for item in lines.split(","):
                        if item == "0":
                            temp_available.append(False)
                        elif item == "1":
                            temp_available.append(True)
                    available_list.append(temp_available)
                    count += 1
                elif count == 1:
                    temp_offset = []
                    for item in lines.split(","):
                        temp = item.split(" ")
                        temp_offset.append([int(temp[0]), int(temp[1])])
                    offset_list.append(temp_offset)
                    count += 1
                elif count == 2:
                    register_time_list.append(float(lines))
        return available_list, offset_list, register_time_list

    def save_roi_images(self, index, last_roi, next_roi):
        folder_name = str(self.video_name[6]) + "_" + \
                         str(self.video_name[8]) + "_" + \
                         str(index).zfill(5)
        project_address = os.getcwd()
        dataset_address = os.path.join(os.path.join(project_address, "datasets"), "used_for_nets")
        if int(self.video_name[8]) < 4:
            dataset_address = os.path.join(dataset_address, "train")
        elif int(self.video_name[8]) < 5:
            dataset_address = os.path.join(dataset_address, "val")
        elif int(self.video_name[8]) == 5:
            dataset_address = os.path.join(dataset_address, "test")
        folder_address = os.path.join(dataset_address, folder_name)
        # self.make_out_dir(folder_address)
        # cv2.imwrite(os.path.join(folder_address, folder_name + "_1.png", last_roi))
        # cv2.imwrite(os.path.join(folder_address, folder_name + "_2.png", next_roi))


# -*- coding: utf-8 -*-
from __future__ import division, print_function, absolute_import

import numpy as np

from alpharotate.utils.pretrain_zoo import PretrainModelZoo
from configs._base_.models.retinanet_r50_fpn import *
from configs._base_.datasets.dota_detection import *
from configs._base_.schedules.schedule_1x import *

# schedule
BATCH_SIZE = 1
GPU_GROUP = "0"
NUM_GPU = len(GPU_GROUP.strip().split(','))
LR = 1e-3
SAVE_WEIGHTS_INTE = 10000
DECAY_STEP = np.array(DECAY_EPOCH, np.int32) * SAVE_WEIGHTS_INTE
MAX_ITERATION = SAVE_WEIGHTS_INTE * MAX_EPOCH
WARM_SETP = int(WARM_EPOCH * SAVE_WEIGHTS_INTE)

# dataset
DATASET_NAME = 'MLT'
IMG_SHORT_SIDE_LEN = 800
IMG_MAX_LENGTH = 1000
CLASS_NUM = 1

# model
pretrain_zoo = PretrainModelZoo()
PRETRAINED_CKPT = pretrain_zoo.pretrain_weight_path(NET_NAME, ROOT_PATH)
TRAINED_CKPT = os.path.join(ROOT_PATH, 'output/trained_weights')

# post-processing
VIS_SCORE = 0.2

VERSION = 'RetinaNet_MLT_KF_1x_20210730'

"""
FLOPs: 472715469;    Trainable params: 32325246

2021-07-31  kf_0.4	55.96%	72.77%	45.46%	41.84%
2021-07-31  kf_0.45	55.59%	76.36%	43.70%	40.52%
2021-07-31  kl_0.35	55.78%	68.31%	47.13%	43.01%
2021-07-31  kf_0.3	54.96%	63.03%	48.72%	44.06%
2021-07-31	kf_0.25	53.16%	56.39%	50.28%	44.99%
"""


#!/usr/bin/python3

import argparse
import itertools

import torchvision.transforms as transforms
from torch.utils.data import DataLoader
from torch.autograd import Variable
from PIL import Image
import torch

from model import Generator, GlobalGenerator2, InceptionV3
# from utils import ReplayBuffer
from utils import LambdaLR
from utils import channel2width
from utils import weights_init_normal
from utils import createNRandompatches
from dataset import UnpairedDepthDataset
import utils_pl
from collections import OrderedDict
import util.util as util
import networks
import numpy as np
import os

parser = argparse.ArgumentParser()
parser.add_argument('--name', type=str, help='name of this experiment')
parser.add_argument('--checkpoints_dir', type=str, default='checkpoints', help='Where checkpoints are saved')
parser.add_argument('--epoch', type=int, default=0, help='starting epoch')
parser.add_argument('--n_epochs', type=int, default=100, help='number of epochs of training')
parser.add_argument('--batchSize', type=int, default=4, help='size of the batches')
parser.add_argument('--cuda', action='store_true', help='use GPU computation', default=True)
parser.add_argument('--n_cpu', type=int, default=8, help='number of cpu threads to use during batch generation')

###loading data
parser.add_argument('--dataroot', type=str, default='datasets/vangogh2photo/', help='photograph directory root directory')
parser.add_argument('--root2', type=str, default='', help='line drawings dataset root directory')
parser.add_argument('--depthroot', type=str, default='', help='dataset of corresponding ground truth depth maps')
parser.add_argument('--feats2Geom_path', type=str, default='checkpoints/feats2Geom/feats2depth.pth',
                                help='path to pretrained features to depth map network')

### architecture and optimizers
parser.add_argument('--lr', type=float, default=0.0002, help='initial learning rate')
parser.add_argument('--momentum', type=float, default=0.9, help='momentum for optimizer')
parser.add_argument('--decay_epoch', type=int, default=50, help='epoch to start linearly decaying the learning rate to 0')
parser.add_argument('--size', type=int, default=256, help='size of the data crop (squared assumed)')
parser.add_argument('--input_nc', type=int, default=3, help='number of channels of input data')
parser.add_argument('--output_nc', type=int, default=1, help='number of channels of output data')
parser.add_argument('--geom_nc', type=int, default=3, help='number of channels of geom data')
parser.add_argument('--ngf', type=int, default=64, help='# of gen filters in first conv layer')
parser.add_argument('--ndf', type=int, default=64, help='# of discrim filters in first conv layer')
parser.add_argument('--netD', type=str, default='basic', help='selects model to use for netD')
parser.add_argument('--n_blocks', type=int, default=3, help='number of resnet blocks for generator')
parser.add_argument('--n_layers_D', type=int, default=3, help='only used if netD==n_layers')
parser.add_argument('--norm', type=str, default='instance', help='instance normalization or batch normalization')
parser.add_argument('--disc_sigmoid', type=int, default=0, help='use sigmoid in disc loss')
parser.add_argument('--every_feat', type=int, default=1, help='use transfer features for recog loss')
parser.add_argument('--finetune_netGeom', type=int, default=1, help='make geometry networks trainable')

### loading from checkpoints
parser.add_argument('--load_pretrain', type=str, default='', help='where to load file if wanted')
parser.add_argument('--continue_train', action='store_true', help='continue training: load the latest model')
parser.add_argument('--which_epoch', type=str, default='latest', help='which epoch to load from if continue_train')

### dataset options
parser.add_argument('--mode', type=str, default='train', help='train, val, test, etc')
parser.add_argument('--load_size', type=int, default=286, help='scale images to this size')
parser.add_argument('--crop_size', type=int, default=256, help='then crop to this size')
parser.add_argument('--aspect_ratio', type=float, default=1.0, help='The ratio width/height. The final height of the load image will be crop_size/aspect_ratio')
parser.add_argument('--max_dataset_size', type=int, default=float("inf"), help='Maximum number of samples allowed per dataset. If the dataset directory contains more than max_dataset_size, only a subset is loaded.')
parser.add_argument('--preprocess', type=str, default='resize_and_crop', help='scaling and cropping of images at load time [resize_and_crop | crop | scale_width | scale_width_and_crop | none]')
parser.add_argument('--no_flip', action='store_true', help='if specified, do not flip the images for data augmentation')

######## loss weights
parser.add_argument("--cond_cycle", type=float, default=0.1, help="weight of the appearance reconstruction loss")
parser.add_argument("--condGAN", type=float, default=1.0, help="weight of the adversarial style loss")
parser.add_argument("--cond_recog", type=float, default=10.0, help="weight of the semantic loss")
parser.add_argument("--condGeom", type=float, default=10.0, help="weight of the geometry style loss")

### geometry loss options
parser.add_argument('--use_geom', type=int, default=1, help='include the geometry loss')
parser.add_argument('--midas', type=int, default=1, help='use midas depth map')

### semantic loss options
parser.add_argument('--N_patches', type=int, default=1, help='number of patches for clip')
parser.add_argument('--patch_size', type=int, default=128, help='patchsize for clip')
parser.add_argument('--num_classes', type=int, default=55, help='number of classes for inception')
parser.add_argument('--cos_clip', type=int, default=0, help='use cosine similarity for CLIP semantic loss')

### save options
parser.add_argument('--save_epoch_freq', type=int, default=1000, help='how often to save the latest model in steps')
parser.add_argument('--slow', type=int, default=0, help='only frequently save netG_A, netGeom')
parser.add_argument('--log_int', type=int, default=50, help='display frequency for tensorboard')

opt = parser.parse_args()
print(opt)

checkpoints_dir = opt.checkpoints_dir
name = opt.name

from util.visualizer2 import Visualizer
tensor2im = util.tensor2imv2

visualizer = Visualizer(checkpoints_dir, name, tf_log=True, isTrain=True)
print('------------------- created visualizer -------------------')

if torch.cuda.is_available() and not opt.cuda:
    print("WARNING: You have a CUDA device, so you should probably run with --cuda")

###### Definition of variables ######
# Networks

netG_A = 0
netG_A = Generator(opt.input_nc, opt.output_nc, opt.n_blocks)
netG_B = Generator(opt.output_nc, opt.input_nc, opt.n_blocks)

if opt.use_geom == 1:

    netGeom = GlobalGenerator2(768, opt.geom_nc, n_downsampling=1, n_UPsampling=3)
    netGeom.load_state_dict(torch.load(opt.feats2Geom_path))
    print("Loading pretrained features to depth network from %s"%opt.feats2Geom_path)
    if opt.finetune_netGeom == 0:
        netGeom.eval()
else:
    opt.finetune_netGeom = 0

D_input_nc_B = opt.output_nc
D_input_nc_A = opt.input_nc

netD_B = networks.define_D(D_input_nc_B, opt.ndf, opt.netD, opt.n_layers_D, opt.norm, use_sigmoid=False)
netD_A = networks.define_D(D_input_nc_A, opt.ndf, opt.netD, opt.n_layers_D, opt.norm, use_sigmoid=False)

device = 'cpu'
if opt.cuda:
    netG_A.cuda()
    netG_B.cuda()
    netD_A.cuda()
    netD_B.cuda()
    if opt.use_geom==1:
        netGeom.cuda()
    device = 'cuda'

### load pretrained inception
net_recog = InceptionV3(opt.num_classes, opt.mode=='test', use_aux=True, pretrain=True, freeze=True, every_feat=opt.every_feat==1)
net_recog.cuda()
net_recog.eval()

import clip
device = "cuda" if torch.cuda.is_available() else "cpu"
clip_model, preprocess = clip.load("ViT-B/32", device=device, jit=False)
clip.model.convert_weights(clip_model)

#### load in progress weights if continue train or load_pretrain
if opt.continue_train:
    netG_A.load_state_dict(torch.load(os.path.join(opt.checkpoints_dir, opt.name, 'netG_A_%s.pth' % opt.which_epoch)))
    netG_B.load_state_dict(torch.load(os.path.join(opt.checkpoints_dir, opt.name, 'netG_B_%s.pth' % opt.which_epoch)))
    netD_A.load_state_dict(torch.load(os.path.join(opt.checkpoints_dir, opt.name, 'netD_A_%s.pth' % opt.which_epoch)))
    netD_B.load_state_dict(torch.load(os.path.join(opt.checkpoints_dir, opt.name, 'netD_B_%s.pth' % opt.which_epoch)))
    if opt.finetune_netGeom == 1:
        netGeom.load_state_dict(torch.load(os.path.join(opt.checkpoints_dir, opt.name, 'netGeom_%s.pth'% opt.which_epoch)))
    print('----------- loaded %s from '%opt.which_epoch + os.path.join(checkpoints_dir, name) + '---------------------- !!')
elif len(opt.load_pretrain) > 0:
    pretrained_path = opt.load_pretrain
    netG_A.load_state_dict(torch.load(os.path.join(pretrained_path, 'netG_A_%s.pth' % opt.which_epoch)))
    netG_B.load_state_dict(torch.load(os.path.join(pretrained_path, 'netG_B_%s.pth' % opt.which_epoch)))
    netD_A.load_state_dict(torch.load(os.path.join(pretrained_path, 'netD_A_%s.pth' % opt.which_epoch)))
    netD_B.load_state_dict(torch.load(os.path.join(pretrained_path, 'netD_B_%s.pth' % opt.which_epoch)))
    if opt.finetune_netGeom == 1:
        netGeom.load_state_dict(torch.load(os.path.join(pretrained_path, 'netGeom_%s.pth'% opt.which_epoch)))
    print('----------- loaded %s from '%opt.which_epoch + ' ' + pretrained_path + '---------------------- !!')
else:
    netG_A.apply(weights_init_normal)
    netG_B.apply(weights_init_normal)
    netD_A.apply(weights_init_normal)
    netD_B.apply(weights_init_normal)

print('----------- loaded networks ---------------------- !!')

# Losses

criterionGAN = networks.GANLoss(use_lsgan=True, target_real_label=1.0, target_fake_label=0.0, reduceme=True).to(device)

criterion_MSE = torch.nn.MSELoss(reduce=True)
criterionCycle = torch.nn.L1Loss()
criterionCycleB = criterionCycle

criterionCLIP = criterion_MSE
if opt.cos_clip == 1:
    criterionCLIP = torch.nn.CosineSimilarity(dim=1, eps=1e-08)

criterionGeom = torch.nn.BCELoss(reduce=True)

############### only use B to A ###########################
optimizer_G_A = torch.optim.Adam(netG_A.parameters(), lr=opt.lr, betas=(0.5, 0.999))
optimizer_G_B = torch.optim.Adam(netG_B.parameters(), lr=opt.lr, betas=(0.5, 0.999))

if (opt.use_geom == 1 and opt.finetune_netGeom==1):
    optimizer_Geom = torch.optim.Adam(netGeom.parameters(), lr=opt.lr, betas=(0.5, 0.999))

optimizer_D_B = torch.optim.Adam(netD_B.parameters(), lr=opt.lr, betas=(0.5, 0.999))
optimizer_D_A = torch.optim.Adam(netD_A.parameters(), lr=opt.lr, betas=(0.5, 0.999))

lr_scheduler_G_A = torch.optim.lr_scheduler.LambdaLR(optimizer_G_A,
                                                   lr_lambda=LambdaLR(opt.n_epochs, opt.epoch, opt.decay_epoch).step)
lr_scheduler_D_B = torch.optim.lr_scheduler.LambdaLR(optimizer_D_B,
                                                     lr_lambda=LambdaLR(opt.n_epochs, opt.epoch, opt.decay_epoch).step)

lr_scheduler_G_B = torch.optim.lr_scheduler.LambdaLR(optimizer_G_B,
                                                   lr_lambda=LambdaLR(opt.n_epochs, opt.epoch, opt.decay_epoch).step)
lr_scheduler_D_A = torch.optim.lr_scheduler.LambdaLR(optimizer_D_A,
                                                     lr_lambda=LambdaLR(opt.n_epochs, opt.epoch, opt.decay_epoch).step)

# Inputs & targets memory allocation
Tensor = torch.cuda.FloatTensor if opt.cuda else torch.Tensorreal_A

# Dataset loader
transforms_r = [transforms.Resize(int(opt.size * 1.12), Image.BICUBIC),
               transforms.RandomCrop(opt.size),
               transforms.ToTensor()]

train_data = UnpairedDepthDataset(opt.dataroot, opt.root2, opt, transforms_r=transforms_r,
                mode=opt.mode, midas=opt.midas>0, depthroot=opt.depthroot)

dataloader = DataLoader(train_data, batch_size=opt.batchSize, shuffle=True, num_workers=opt.n_cpu, drop_last=True)

print('---------------- loaded %d images ----------------' % len(train_data))

###### Training ######
for epoch in range(opt.epoch, opt.n_epochs):
    for i, batch in enumerate(dataloader):
        total_steps = epoch*len(dataloader) + i

        img_r  = Variable(batch['r']).cuda()
        img_depth  = Variable(batch['depth']).cuda()

        real_A = img_r
        labels = Variable(batch['label']).cuda()

        real_B = 0
        real_B  = Variable(batch['line']).cuda()

        recover_geom = img_depth
        batch_size = real_A.size()[0]

        condGAN = opt.condGAN
        cond_recog = opt.cond_recog
        cond_cycle = opt.cond_cycle

        #################### Generator ####################

        fake_B = netG_A(real_A) # G_A(A)
        rec_A = netG_B(fake_B)   # G_B(G_A(A))

        fake_A = netG_B(real_B)  # G_B(B)
        rec_B = netG_A(fake_A) # G_A(G_B(B))

        loss_cycle_Geom = 0
        if opt.use_geom == 1:
            geom_input = fake_B
            if geom_input.size()[1] == 1:
                geom_input = geom_input.repeat(1, 3, 1, 1)
            _, geom_input = net_recog(geom_input)

            pred_geom = netGeom(geom_input)
            pred_geom = (pred_geom+1)/2.0 ###[-1, 1] ---> [0, 1]

            loss_cycle_Geom = criterionGeom(pred_geom, recover_geom)

        ########## loss A Reconstruction ##########

        loss_G_A = criterionGAN(netD_A(fake_A), True)

        # GAN loss D_B(G_B(B))
        loss_G_B = 0
        pred_fake_GAN = netD_B(fake_B)
        loss_G_B = criterionGAN(netD_B(fake_B), True)

        # Forward cycle loss || G_B(G_A(A)) - A||
        loss_cycle_A = criterionCycle(rec_A, real_A)
        loss_cycle_B = criterionCycleB(rec_B, real_B)
        # combined loss and calculate gradients

        loss_GAN = loss_G_A + loss_G_B
        loss_RC = loss_cycle_A + loss_cycle_B

        loss_G = cond_cycle*loss_RC + condGAN*loss_GAN
        loss_G += opt.condGeom*loss_cycle_Geom

        ### semantic loss
        loss_recog = 0

        # renormalize mean=(0.48145466, 0.4578275, 0.40821073), std=(0.26862954, 0.26130258, 0.27577711)
        recog_real = real_A
        recog_real0 = (recog_real[:, 0, :, :].unsqueeze(1) - 0.48145466) / 0.26862954
        recog_real1 = (recog_real[:, 1, :, :].unsqueeze(1) - 0.4578275) / 0.26130258
        recog_real2 = (recog_real[:, 2, :, :].unsqueeze(1) - 0.40821073) / 0.27577711
        recog_real = torch.cat([recog_real0, recog_real1, recog_real2], dim=1)

        line_input = fake_B
        if opt.output_nc == 1:
            line_input_channel0 = (line_input - 0.48145466) / 0.26862954
            line_input_channel1 = (line_input - 0.4578275) / 0.26130258
            line_input_channel2 = (line_input - 0.40821073) / 0.27577711
            line_input = torch.cat([line_input_channel0, line_input_channel1, line_input_channel2], dim=1)

        patches_r = [torch.nn.functional.interpolate(recog_real, size=224)]  #The resize operation on tensor.
        patches_l = [torch.nn.functional.interpolate(line_input, size=224)]

        ## patch based clip loss
        if opt.N_patches > 1:
            patches_r2, patches_l2 = createNRandompatches(recog_real, line_input, opt.N_patches, opt.patch_size)
            patches_r += patches_r2
            patches_l += patches_l2

        loss_recog = 0
        for patchnum in range(len(patches_r)):

            real_patch = patches_r[patchnum]
            line_patch = patches_l[patchnum]

            feats_r = clip_model.encode_image(real_patch).detach()
            feats_line = clip_model.encode_image(line_patch)

            myloss_recog = criterionCLIP(feats_line, feats_r.detach())
            if opt.cos_clip == 1:
                myloss_recog = 1.0 - loss_recog
                myloss_recog = torch.mean(loss_recog)

            patch_factor = (1.0 / float(opt.N_patches))
            if patchnum == 0:
                patch_factor = 1.0
            loss_recog += patch_factor*myloss_recog

        loss_G += cond_recog* loss_recog

        optimizer_G_A.zero_grad()
        optimizer_G_B.zero_grad()
        if opt.finetune_netGeom == 1:
            optimizer_Geom.zero_grad()

        loss_G.backward()

        optimizer_G_A.step()
        optimizer_G_B.step()
        if opt.finetune_netGeom == 1:
            optimizer_Geom.step()

        ##########  Discriminator A ##########

        # Fake loss
        pred_fake_A = netD_A(fake_A.detach())
        loss_D_A_fake = criterionGAN(pred_fake_A, False)

        # Real loss

        pred_real_A = netD_A(real_A)
        loss_D_A_real = criterionGAN(pred_real_A, True)

        # Total loss
        loss_D_A = torch.mean(condGAN * (loss_D_A_real + loss_D_A_fake) ) * 0.5

        optimizer_D_A.zero_grad()
        loss_D_A.backward()
        optimizer_D_A.step()

        ##########  Discriminator B ##########

        # Fake loss
        pred_fake_B = netD_B(fake_B.detach())
        loss_D_B_fake = criterionGAN(pred_fake_B, False)

        # Real loss

        pred_real_B = netD_B(real_B)
        loss_D_B_real = criterionGAN(pred_real_B, True)

        # Total loss
        loss_D_B = torch.mean(condGAN * (loss_D_B_real + loss_D_B_fake) ) * 0.5
        optimizer_D_B.zero_grad()
        loss_D_B.backward()
        optimizer_D_B.step()

        # Progress report
        if (i+1)%opt.log_int==0:

            errors = {}
            errors['total_G'] = loss_G.item() if not isinstance(loss_G, (int,float)) else loss_G
            errors['loss_RC'] = torch.mean(loss_RC) if not isinstance(loss_RC, (int,float)) else loss_RC
            errors['loss_cycle_Geom'] = torch.mean(loss_cycle_Geom) if not isinstance(loss_cycle_Geom, (int,float)) else loss_cycle_Geom
            errors['loss_GAN'] = torch.mean(loss_GAN) if not isinstance(loss_GAN, (int,float)) else loss_GANB
            errors['loss_D_B'] = loss_D_B.item() if not isinstance(loss_D_B, (int,float)) else loss_D_B
            errors['loss_D_A'] = loss_D_A.item() if not isinstance(loss_D_A, (int,float)) else loss_D_A
            errors['loss_recog'] = torch.mean(loss_recog) if not isinstance(loss_recog, (int,float)) else loss_recog

            visualizer.print_current_errors(epoch, total_steps, errors, 0)
            visualizer.plot_current_errors(errors, total_steps)

            with torch.no_grad():

                input_img = channel2width(real_A)
                if opt.use_geom == 1:
                    pred_geom = channel2width(pred_geom)
                    input_img = torch.cat([input_img, channel2width(recover_geom)], dim=3)

                input_img_fake = channel2width(fake_A)
                rec_A = channel2width(rec_A)

                show_real_B = real_B

                visuals = OrderedDict([('real_A', tensor2im(input_img.data[0])),
                                           ('real_B', tensor2im(show_real_B.data[0])),
                                           ('fake_A', tensor2im(input_img_fake.data[0])),
                                           ('rec_A', tensor2im(rec_A.data[0])),
                                           ('fake_B', tensor2im(fake_B.data[0]))])

                if opt.use_geom == 1:
                    visuals['pred_geom'] = tensor2im(pred_geom.data[0])

                visualizer.display_current_results(visuals, total_steps, epoch)

    # Update learning rates
    lr_scheduler_G_A.step()
    lr_scheduler_G_B.step()
    lr_scheduler_D_A.step()
    lr_scheduler_D_B.step()

    # Save models checkpoints
    # torch.save(netG_A2B.state_dict(), 'output/netG_A2B.pth')
    if (epoch+1) % opt.save_epoch_freq == 0:
        torch.save(netG_A.state_dict(), os.path.join(opt.checkpoints_dir, name, 'netG_A_%02d.pth'%(epoch)))
        if opt.finetune_netGeom == 1:
            torch.save(netGeom.state_dict(), os.path.join(opt.checkpoints_dir, name, 'netGeom_%02d.pth'%(epoch)))
        if opt.slow == 0:
            torch.save(netG_B.state_dict(), os.path.join(opt.checkpoints_dir, name, 'netG_B_%02d.pth'%(epoch)))
            torch.save(netD_A.state_dict(), os.path.join(opt.checkpoints_dir, name, 'netD_A_%02d.pth'%(epoch)))
            torch.save(netD_B.state_dict(), os.path.join(opt.checkpoints_dir, name, 'netD_B_%02d.pth'%(epoch)))

    torch.save(netG_A.state_dict(), os.path.join(opt.checkpoints_dir, name, 'netG_A_latest.pth'))
    torch.save(netG_B.state_dict(), os.path.join(opt.checkpoints_dir, name, 'netG_B_latest.pth'))
    torch.save(netD_B.state_dict(), os.path.join(opt.checkpoints_dir, name, 'netD_B_latest.pth'))
    torch.save(netD_A.state_dict(), os.path.join(opt.checkpoints_dir, name, 'netD_A_latest.pth'))
    if opt.finetune_netGeom == 1:
        torch.save(netGeom.state_dict(), os.path.join(opt.checkpoints_dir, name, 'netGeom_latest.pth'))

###################################


import io
import os
import sys
import re
import timeit
import random

import tensorflow as tf
import editdistance
import numpy as np
from scipy import exp
from scipy.special import lambertw

from ai.datasets import ALLDATASET
from ai.models import Seq2Seq

# INPUT FILES
tf.app.flags.DEFINE_string('train_input', None, "Path to train input file")
tf.app.flags.DEFINE_string('train_output', None, "Path to train output file")
tf.app.flags.DEFINE_string('dev_input', None, "Path to dev input file")
tf.app.flags.DEFINE_string('dev_output', None, "Path to dev output file")

# HYPERPARAMETERS
tf.app.flags.DEFINE_float('lr', 0.0001, "Initial learning rate.")
tf.app.flags.DEFINE_integer('batch_size', 2048, "Batch size.")
tf.app.flags.DEFINE_integer('embedding_size', 256, "Embedding dimensionality.")
tf.app.flags.DEFINE_integer('hidden_size', 256, "Number of hidden units.")
tf.app.flags.DEFINE_integer('rnn_layers', 2, "Number of RNN layers.")
tf.app.flags.DEFINE_boolean('bidirectional_encoder', True, "Whether to use a"
                            " bidirectional RNN in the encoder's 1st layer.")
tf.app.flags.DEFINE_string('bidirectional_mode', 'add', "Set to 'add',"
                           " 'concat' or 'project'.")
tf.app.flags.DEFINE_boolean('use_lstm', False, "Set to False to use GRUs.")
tf.app.flags.DEFINE_string('attention', 'luong', "'bahdanau' or 'luong'"
                           " (default is 'luong').")
tf.app.flags.DEFINE_float('dropout', .9, "Keep probability for dropout on the"
                          "RNNs' non-recurrent connections.")
tf.app.flags.DEFINE_float('max_grad_norm', 11.5, "Clip gradients to this norm.")
tf.app.flags.DEFINE_integer('beam_size', 5, "Beam search size.")
tf.app.flags.DEFINE_float('initial_p_sample', .35, "Initial decoder sampling"
                          " probability (0=ground truth, 1=use predictions).")
tf.app.flags.DEFINE_float('final_p_sample', .35, "Final decoder sampling"
                          " probability (0=ground truth, 1=use predictions).")
tf.app.flags.DEFINE_integer('epochs_p_sample', 20, "Duration in epochs of"
                            " schedule sampling (determines rate of change).")
tf.app.flags.DEFINE_boolean('linear_p_sample', True, "False = sigmoid decay.")
tf.app.flags.DEFINE_integer('parse_repeated', 0, "Set to > 1 to compress"
                            " contiguous patterns in the data pipeline.")
tf.app.flags.DEFINE_float('epsilon', 1e-8, "Denominator constant.")
tf.app.flags.DEFINE_float('beta1', .9, "First order moment decay.")
tf.app.flags.DEFINE_float('beta2', .999, "Second order moment decay.")

# CONFIG
tf.app.flags.DEFINE_integer('max_sentence_length', 110, "Max. word length of"
                            " training examples (both inputs and labels).")
tf.app.flags.DEFINE_integer('num_steps_per_eval', 50, "Number of steps to wait"
                            " before running the graph with the dev set.")
tf.app.flags.DEFINE_integer('max_epochs', 40, "Number of epochs to run"
                            " (0 = no limit).")
tf.app.flags.DEFINE_string('predict_input_file', None, "Path to prediction input file")
tf.app.flags.DEFINE_string('predict_output_file', None, "Name of the output file with decoding results.")
tf.app.flags.DEFINE_boolean('restore', True, "Whether to restore the model.")
tf.app.flags.DEFINE_string('model_output_dir', None, "Name of the output directory.")

FLAGS = tf.app.flags.FLAGS

def untokenize_batch(dataset, id_batch):
  """Return the UTF-8 sequences of the given batch of ids."""
  return [dataset.untokenize(dataset.clean(s)) for s in id_batch]

def levenshtein(proposed, gold, normalize=False):
  """Return the normalized Levenshtein distance of the given strings."""
  lev_densities = []
  for x, y in zip(proposed, gold):
    score = editdistance.eval(x, y)
    if normalize:
      score /= len(y)
    lev_densities.append(score)
  return sum(lev_densities) / len(lev_densities)

def train():
  """Run a loop that continuously trains the model."""
  print("Building dynamic character-level ALLDATASET data...", flush=True)
  dataset = ALLDATASET(
    train_input=FLAGS.train_input, train_output=FLAGS.train_output,
    dev_input=FLAGS.dev_input, dev_output=FLAGS.dev_output,
    predict_input_file=FLAGS.predict_input_file,
    parse_repeated=FLAGS.parse_repeated,
    shuffle=True, max_input_length=FLAGS.max_sentence_length,
    max_label_length=FLAGS.max_sentence_length)

  print("Building computational graph...", flush=True)
  graph = tf.Graph()

  with graph.as_default():

    tf.set_random_seed(1)
    random.seed(1)
    np.random.seed(1)

    # During training we use beam width 1. There are lots of complications on
    # the implementation, e.g. only tiling during inference.
    m = Seq2Seq(
      num_types=dataset.num_types(),
      max_encoder_length=FLAGS.max_sentence_length,
      max_decoder_length=FLAGS.max_sentence_length,
      pad_id=dataset.type_to_ix['_PAD'],
      eos_id=dataset.type_to_ix['_EOS'],
      go_id=dataset.type_to_ix['_GO'],
      space_id=dataset.type_to_ix[(' ',)],
      ix_to_type=dataset.ix_to_type,
      batch_size=FLAGS.batch_size, embedding_size=FLAGS.embedding_size,
      hidden_size=FLAGS.hidden_size, rnn_layers=FLAGS.rnn_layers,
      bidirectional_encoder=FLAGS.bidirectional_encoder,
      bidirectional_mode=FLAGS.bidirectional_mode,
      use_lstm=FLAGS.use_lstm, attention=FLAGS.attention,
      dropout=FLAGS.dropout, max_grad_norm=FLAGS.max_grad_norm, beam_size=1,
      epsilon=FLAGS.epsilon, beta1=FLAGS.beta1, beta2=FLAGS.beta2,
      restore=FLAGS.restore, model_output_dir=FLAGS.model_output_dir)

  # Allow TensorFlow to resort back to CPU when we try to set an operation to
  # a GPU where there's only a CPU implementation, rather than crashing.
  sess_config = tf.ConfigProto(allow_soft_placement=True)

  with tf.Session(graph=graph, config=sess_config) as sess:
    print("Initializing or restoring model...", flush=True)
    m.start()

    # If the model was not restored, initialize the variable hyperparameters.
    if sess.run(m.lr) == 0:
      sess.run(tf.assign(m.lr, FLAGS.lr))
    if sess.run(m.p_sample) == 0:
      sess.run(tf.assign(m.p_sample, FLAGS.initial_p_sample))

    # Get the number of epochs that have passed (easier by getting batches now)
    step = m.global_step.eval()
    batches = dataset.get_train_batches(m.batch_size)
    epoch = step // len(batches)

    # Scheduled sampling decay
    i = FLAGS.initial_p_sample
    f = FLAGS.final_p_sample
    # The stopping point is based on the max epochs
    total_train_steps = len(batches) * FLAGS.epochs_p_sample
    if i != f and not FLAGS.linear_p_sample:
      k = total_train_steps / (float(lambertw(total_train_steps / 2)) * 2)
      expk = float(exp(-total_train_steps / k))
      delta_f = (f - i) * (1 + k) * (1 + k * expk) / (k - k * expk) - f
      delta_i = (f + delta_f) / (1 + k)

    while not FLAGS.max_epochs or epoch <= FLAGS.max_epochs:
      print("=====EPOCH {}=====".format(epoch), flush=True)
      while step < (epoch + 1) * len(batches):
        step = m.global_step.eval()

        # Scheduled sampling decay
        if i != f:
          # Linear decay
          if FLAGS.linear_p_sample:
            p = min(f, i + step * (f - i) / total_train_steps)
          # Inverse sigmoid decay
          else:
            expk = float(exp(-step / k))
            p = min(f, i - delta_i + (f + delta_f) / (1 + k * expk))

          sess.run(tf.assign(m.p_sample, p))

        # Gradient descent and backprop
        train_inputs, train_labels = zip(*batches[step % len(batches)])
        train_fd = {m.inputs: train_inputs, m.labels: train_labels}

        # Wrap into function to measure running time
        def train_step():
          sess.run(m.train_step, feed_dict=train_fd)

        print("Global step {0} ({1}s)".format(
          step, timeit.timeit(train_step, number=1)), flush=True)

        if step % FLAGS.num_steps_per_eval == 0:
          valid_inputs, valid_labels = dataset.get_valid_batch(m.batch_size)
          valid_fd = {m.inputs: valid_inputs, m.labels: valid_labels}

          # Run training and validation perplexity and samples

          lr, train_ppx, train_output, p_sample, train_ppx_summ = sess.run([
            m.lr,
            m.perplexity,
            m.output,
            m.p_sample,
            m.perplexity_summary,
          ], feed_dict=train_fd)

          valid_ppx, valid_output, infer_output, valid_ppx_summ = sess.run([
            m.perplexity,
            m.output,
            m.generative_output,
            m.perplexity_summary,
          ], feed_dict=valid_fd)

          # Convert data to UTF-8 strings for evaluation and display
          valid_inputs = untokenize_batch(dataset, valid_inputs)
          valid_labels = untokenize_batch(dataset, valid_labels)
          valid_output = untokenize_batch(dataset, valid_output)
          infer_output = untokenize_batch(dataset, infer_output)

          # Run evaluation metrics
          lev = levenshtein(infer_output, valid_labels)
          lev_density = levenshtein(infer_output, valid_labels, normalize=True)

          lev_summ = sess.run(
            m.lev_summary, feed_dict={m.lev: lev})
          lev_density_summ = sess.run(
            m.lev_density_summary, feed_dict={m.lev_density: lev_density})

          # Write summaries to TensorBoard
          m.train_writer.add_summary(train_ppx_summ, global_step=step)
          m.valid_writer.add_summary(valid_ppx_summ, global_step=step)
          m.valid_writer.add_summary(lev_summ, global_step=step)
          m.valid_writer.add_summary(lev_density_summ, global_step=step)

          # Display results to stdout
          print("  lr:", lr)
          print("  p_sample:", p_sample)
          print("  train_ppx:", train_ppx)
          print("  valid_ppx:", valid_ppx)
          print("  lev:", lev)
          print("  lev_density:", lev_density)
          print("Input:")
          print(valid_inputs[0])
          print("Target:")
          print(valid_labels[0])
          print("Output with ground truth:")
          print(valid_output[0])
          print("Greedily decoded output:")
          print(infer_output[0], flush=True)

      # Epoch about to be done - save, reshuffle the data and get new batches
      print("Saving model...")
      m.save()
      print("Model saved. Resuming training...", flush=True)
      batches = dataset.get_train_batches(m.batch_size)
      epoch += 1

def decode():
  """Run a blind test on the file with path given by the `decode` flag."""
  with io.open(FLAGS.predict_input_file, encoding='utf-8') as test_file:
    lines = test_file.readlines()
    # Get the largest sentence length to set an upper bound to the decoder.
    max_length = FLAGS.max_sentence_length
    # max_length = max([len(line) for line in lines])

  print("Building dynamic character-level ALLDATASET data...", flush=True)
  dataset = ALLDATASET(
    train_input=FLAGS.train_input, train_output=FLAGS.train_output,
    dev_input=FLAGS.dev_input, dev_output=FLAGS.dev_output,
    predict_input_file=FLAGS.predict_input_file,
    parse_repeated=FLAGS.parse_repeated,
    max_input_length=max_length, max_label_length=max_length)

  print("Building computational graph...", flush=True)
  graph = tf.Graph()
  with graph.as_default():

    tf.set_random_seed(1)
    random.seed(1)
    np.random.seed(1)

    m = Seq2Seq(
      num_types=dataset.num_types(),
      max_encoder_length=max_length, max_decoder_length=max_length,
      pad_id=dataset.type_to_ix['_PAD'],
      eos_id=dataset.type_to_ix['_EOS'],
      go_id=dataset.type_to_ix['_GO'],
      space_id=dataset.type_to_ix[(' ',)],
      ix_to_type=dataset.ix_to_type,
      batch_size=1, embedding_size=FLAGS.embedding_size,
      hidden_size=FLAGS.hidden_size, rnn_layers=FLAGS.rnn_layers,
      bidirectional_encoder=FLAGS.bidirectional_encoder,
      bidirectional_mode=FLAGS.bidirectional_mode,
      use_lstm=FLAGS.use_lstm, attention=FLAGS.attention,
      beam_size=FLAGS.beam_size, restore=True, model_output_dir=FLAGS.model_output_dir)

  with tf.Session(graph=graph) as sess:
    print("Restoring model...", flush=True)
    m.start()
    print(
      "Restored model (global step {})".format(m.global_step.eval()),
      flush=True)
    with io.open(FLAGS.predict_output_file, 'w', encoding='utf-8') as output_file:
      for line in lines:
        # if len(line) > max_length:
        #   continue
        number_of_chars = len(line)
        completely_divisble = number_of_chars % FLAGS.max_sentence_length == 0

        if number_of_chars < FLAGS.max_sentence_length:
          parts = [line]
        else:
          parts = []
          count = 0
          last_word_end_index = 0

          line_copy = line
          while len(line_copy) != 0 and count < len(line_copy):
            if count == FLAGS.max_sentence_length:
              if last_word_end_index == 0:
                parts.append(line_copy[: count])
                line_copy = line_copy[count:]
              else:
                parts.append(line_copy[: last_word_end_index])
                line_copy = line_copy[last_word_end_index:]

                last_word_end_index = 0
                count = 0

            if line_copy[count] == " ":
              last_word_end_index = count

            count += 1

          if not completely_divisble:
            parts.append(line_copy)

        result = ""
        for part in parts:
          ids = dataset.tokenize(part)
          while len(ids) < max_length:
            ids.append(dataset.type_to_ix['_PAD'])
          outputs = sess.run(m.generative_output, feed_dict={m.inputs: [ids]})
          top_line = untokenize_batch(dataset, outputs)[0]
          # Sequences of text will only be repeated up to 5 times.
          top_line = re.sub(r'(.+?)\1{5,}', lambda m: m.group(1) * 5, top_line)
          result += top_line
        output_file.write(result + '\n')
        print("PREDICTION:", top_line, flush=True)
        print()

def main(_):
  """Callee of `tf.app.run` the method."""
  if not FLAGS.model_output_dir:
    raise ValueError(
      "Undefined model output directory. Perhaps you forgot to set the --model_output_dir flag?")

  if FLAGS.predict_input_file:
    decode()
  else:
    train()

if __name__ == '__main__':
  tf.app.run()


#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
Gym trading env
https://github.com/openai/gym/blob/master/gym/envs/classic_control/cartpole.py
1. obs <- reset()      # env
2. action <- pi(obs)    # agent
3. news_obs <- step(action)      # env
repeat 2, and 3 for interactions between agent and env
"""
import os
import numpy as np
import pandas as pd
import gym
from datetime import datetime, date
import logging

from .event import EventType
from .event.backtest_event_engine import BacktestEventEngine
from .data.backtest_data_feed import BacktestDataFeed
from .data.data_board import DataBoard
from .brokerage.backtest_brokerage import BacktestBrokerage
from .position.position_manager import PositionManager
from .performance.performance_manager import PerformanceManager
from .risk.risk_manager import PassThroughRiskManager

_logger = logging.getLogger(__name__)

class BacktestGymEngine(gym.Env):
    """
    Description:
        backtest gym engine
        it doesn't normalize; and expects a normalization layer
    Observation:
        Type: Box(lookback_window, n_assets*5+2)
        lookback_window x (n_assets*(ohlcv) + cash+npv)
        TODO: append trades, commissions, standing orders, etc
        TODO: stop/limit orders
    Actions:
        Type: Box(n_assets + 1)
        portfolio weights [w1,w2...w_k, cash_weight], add up to one
    Reward:
        cumulative pnl in run_window
    Starting State:
        random timestamp between start_date and (end_date - run_window)
    Episode Termination:
        after predefined window
        If broke, no orders will send
    """
    def __init__(self, n_assets, lookback_window=15, run_window=252*2, start_date=None, end_date=None):
        self._current_time = None
        self._start_date = start_date
        self._end_date = end_date
        self.multiplier_dict = {}
        self._data_feed = BacktestDataFeed(self._start_date, self._end_date)
        self._data_board = DataBoard()
        self._performance_manager = PerformanceManager(self.multiplier_dict)
        self._position_manager = PositionManager('Global')
        self._position_manager.set_multiplier(self.multiplier_dict)
        self._risk_manager = PassThroughRiskManager()
        self._strategy = None           # no strategy; strategy is to be learned

        self._n_assets = n_assets
        self.action_space = gym.spaces.Box(low=0.0, high=1.0, shape=(n_assets + 1,), dtype=np.float32)

        self._lookback_window = lookback_window
        self._run_window = run_window
        self.observation_space = gym.spaces.Box(low=-np.inf, high=np.inf,
                                                shape=(self._lookback_window, self._n_assets*5+2), dtype=np.float32)

        self._setup()

    def set_capital(self, capital):
        self._position_manager.set_capital(capital)

    def set_multiplier(self, multiplier_dict):
        self.multiplier_dict.update(multiplier_dict)

    def add_data(self, data_key, data_source, watch=True):
        """
        Add data for backtest
        :param data_key: AAPL or CL; if it is followed by number, assumed to be multiplier
        :param data_source:  dataframe, datetimeindex
        :param watch: track position or not
        :return:
        """
        keys = data_key.split(' ')
        if keys[-1].isdigit():       # multiplier
            data_key = ' '.join(keys[:-1])
            self.multiplier_dict[data_key] = int(keys[-1])

        self._data_feed.set_data_source(data_source)          # get iter(datetimeindex)
        self._data_board.initialize_hist_data(data_key, data_source)
        if watch:
            self._performance_manager.add_watch(data_key, data_source)

    def _setup(self):
        """
        Tis can be initialzied before data is loaded
        because it doesn't subscribe
        """
        ## 2. event engine
        self._events_engine = BacktestEventEngine(self._data_feed)

        ## 3. brokerage
        self._backtest_brokerage = BacktestBrokerage(
            self._events_engine, self._data_board
        )

        ## 6. wire up event handlers
        self._events_engine.register_handler(EventType.TICK, self._tick_event_handler)
        # to be consistent with current live, order is placed directly
        # self._events_engine.register_handler(EventType.ORDER, self._order_event_handler)
        self._events_engine.register_handler(EventType.FILL, self._fill_event_handler)

    def _tick_event_handler(self, tick_event):
        self._current_time = tick_event.timestamp

        # performance updates after one step run
        self._position_manager.mark_to_market(tick_event.timestamp, tick_event.full_symbol, tick_event.price, self._data_board)
        self._data_board.on_tick(tick_event)
        self._strategy.on_tick(tick_event)

    def _order_event_handler(self, order_event):
        """
        backtest doesn't send order_event back to strategy. It fills directly and becoems fill_event
        """
        self._backtest_brokerage.place_order(order_event)

    def _fill_event_handler(self, fill_event):
        self._position_manager.on_fill(fill_event)
        self._performance_manager.on_fill(fill_event)
        self._strategy.on_fill(fill_event)

    def reset(self):
        ## reset performance manager and portfolio manager
        self._performance_manager.reset()
        self._position_manager.reset()

        ## reset iterator randomly
        self._data_feed.subscribe_market_data()

    def step(self, action):
        """
        Run backtest
        """
        self._events_engine.run(1)
        # explicitly update performance; now orders are filled
        self._performance_manager.update_performance(self._current_time, self._position_manager, self._data_board)

        done = self._current_time is None
        state = None
        return state

    def render(self, mode='human'):
        pass

    def close(self):
        pass


import argparse
import numpy as np
import pandas as pd
import os
import cv2
from tqdm import tqdm
import torch.nn as nn
from torch import optim
import torch.nn.functional as F
from torchvision import transforms
from torch.utils.data import Dataset, DataLoader
import torch
from apex import amp
from efficientnet_pytorch import EfficientNet
import random
from sklearn.metrics import roc_auc_score
import pickle
import pydicom

class AverageMeter(object):
    """Computes and stores the average and current value"""
    def __init__(self):
        self.reset()
    def reset(self):
        self.val = 0
        self.avg = 0
        self.sum = 0
        self.count = 0
    def update(self, val, n=1):
        self.val = val
        self.sum += val * n
        self.count += n
        self.avg = self.sum / self.count

def window(x, WL=50, WW=350):
    upper, lower = WL+WW//2, WL-WW//2
    x = np.clip(x, lower, upper)
    x = x - np.min(x)
    x = x / np.max(x)
    return x

class PEDataset(Dataset):
    def __init__(self, image_dict, bbox_dict, image_list, target_size):
        self.image_dict=image_dict
        self.bbox_dict=bbox_dict
        self.image_list=image_list
        self.target_size=target_size
    def __len__(self):
        return len(self.image_list)
    def __getitem__(self,index):
        study_id = self.image_dict[self.image_list[index]]['series_id'].split('_')[0]
        series_id = self.image_dict[self.image_list[index]]['series_id'].split('_')[1]
        data = pydicom.dcmread('../../../input/train/'+study_id+'/'+series_id+'/'+self.image_list[index]+'.dcm')
        x = data.pixel_array.astype(np.float32)
        x = x*data.RescaleSlope+data.RescaleIntercept
        x1 = window(x, WL=100, WW=700)
        x = np.zeros((x1.shape[0], x1.shape[1], 3), dtype=np.float32)
        x[:,:,0] = x1
        x[:,:,1] = x1
        x[:,:,2] = x1
        x = cv2.resize(x, (self.target_size,self.target_size))
        x = x.transpose(2, 0, 1)
        bboxes = self.bbox_dict[self.image_list[index]]
        y = torch.from_numpy(np.array(bboxes))
        return x, y

class efficientnet(nn.Module):
    def __init__(self ):
        super().__init__()
        self.net = EfficientNet.from_pretrained('efficientnet-b0')
        in_features = self.net._fc.in_features
        self.last_linear = nn.Linear(in_features, 4)
    def forward(self, x):
        x = self.net.extract_features(x)
        x = self.net._avg_pooling(x)
        x = x.view(x.size(0), -1)
        x = self.last_linear(x)
        return x

def main():

    # checkpoint list
    checkpoint_list = [
                       'epoch20',
                       'epoch21',
                       'epoch22',
                       'epoch23',
                       'epoch24',
                       'epoch25',
                       'epoch26',
                       'epoch27',
                       'epoch28',
                       'epoch29',
                      ]

    # prepare input
    import pickle
    with open('../../process_input/split2/series_list_valid.pickle', 'rb') as f:
        series_list_valid = pickle.load(f)
    with open('../../process_input/split2/series_dict.pickle', 'rb') as f:
        series_dict = pickle.load(f)
    with open('../../process_input/split2/image_dict.pickle', 'rb') as f:
        image_dict = pickle.load(f)
    df = pd.read_csv('../lung_bbox.csv')
    bbox_image_id_list = df['Image'].values
    bbox_Xmin_list = df['Xmin'].values
    bbox_Ymin_list = df['Ymin'].values
    bbox_Xmax_list = df['Xmax'].values
    bbox_Ymax_list = df['Ymax'].values
    bbox_dict = {}
    for i in range(len(bbox_image_id_list)):
        bbox_dict[bbox_image_id_list[i]] = [max(0.0, bbox_Xmin_list[i]), max(0.0, bbox_Ymin_list[i]), min(1.0, bbox_Xmax_list[i]), min(1.0, bbox_Ymax_list[i])]
    image_list_valid = []
    for series_id in series_list_valid:
        sorted_image_list = series_dict[series_id]['sorted_image_list']
        num_image = len(sorted_image_list)
        selected_idx = [int(0.2*num_image), int(0.3*num_image), int(0.4*num_image), int(0.5*num_image)]
        image_list_valid.append(sorted_image_list[selected_idx[0]])
        image_list_valid.append(sorted_image_list[selected_idx[1]])
        image_list_valid.append(sorted_image_list[selected_idx[2]])
        image_list_valid.append(sorted_image_list[selected_idx[3]])
    print(len(image_list_valid))

    # hyperparameters
    batch_size = 256
    image_size = 512
    criterion = nn.L1Loss().cuda()

    # start validation
    for ckp in checkpoint_list:

        # build model
        model = efficientnet()
        model.load_state_dict(torch.load('weights/'+ckp+'_polyak'))
        model = model.cuda()
        optimizer = optim.Adam(model.parameters(), lr=0.0001)
        model, optimizer = amp.initialize(model, optimizer, opt_level="O1",verbosity=0)
        if torch.cuda.device_count() > 1:
            model = torch.nn.DataParallel(model)
        model.eval()

        pred_bbox = np.zeros((len(image_list_valid),4),dtype=np.float32)

        # iterator for validation
        datagen = PEDataset(image_dict=image_dict, bbox_dict=bbox_dict, image_list=image_list_valid, target_size=image_size)
        generator = DataLoader(dataset=datagen, batch_size=batch_size, shuffle=False, num_workers=8, pin_memory=True)

        losses = AverageMeter()
        for i, (images, labels) in enumerate(generator):
            with torch.no_grad():
                start = i*batch_size
                end = start+batch_size
                if i == len(generator)-1:
                    end = len(generator.dataset)
                images = images.cuda()
                labels = labels.float().cuda()
                logits = model(images)
                loss = criterion(logits,labels)
                losses.update(loss.item(), images.size(0))
                pred_bbox[start:end] = np.squeeze(logits.cpu().data.numpy())

        print("checkpoint {} ...".format(ckp))
        print('loss:{}'.format(losses.avg), flush=True)
        print()

if __name__ == "__main__":
    main()


#!/usr/bin/env python3
import os
import gym
import ptan
import argparse
import numpy as np
from tensorboardX import SummaryWriter

import torch
import torch.nn.functional as F
import torch.optim as optim

from lib import common, i2a

LEARNING_RATE = 5e-4
NUM_ENVS = 16
BATCH_SIZE = 64
SAVE_EVERY_BATCH = 1000
OBS_WEIGHT = 10.0
REWARD_WEIGHT = 1.0

def get_obs_diff(prev_obs, cur_obs):
    prev = np.array(prev_obs)[-1]
    cur = np.array(cur_obs)[-1]
    prev = prev.astype(np.float32) / 255.0
    cur = cur.astype(np.float32) / 255.0
    return cur - prev

def iterate_batches(envs, net, device="cpu"):
    act_selector = ptan.actions.ProbabilityActionSelector()
    mb_obs = np.zeros((BATCH_SIZE, ) + common.IMG_SHAPE, dtype=np.uint8)
    mb_obs_next = np.zeros((BATCH_SIZE, ) + i2a.EM_OUT_SHAPE, dtype=np.float32)
    mb_actions = np.zeros((BATCH_SIZE, ), dtype=np.int32)
    mb_rewards = np.zeros((BATCH_SIZE, ), dtype=np.float32)
    obs = [e.reset() for e in envs]
    total_reward = [0.0] * NUM_ENVS
    total_steps = [0] * NUM_ENVS
    batch_idx = 0
    done_rewards = []
    done_steps = []

    while True:
        obs_v = ptan.agent.default_states_preprocessor(obs).to(device)
        logits_v, values_v = net(obs_v)
        probs_v = F.softmax(logits_v, dim=1)
        probs = probs_v.data.cpu().numpy()
        actions = act_selector(probs)

        for e_idx, e in enumerate(envs):
            o, r, done, _ = e.step(actions[e_idx])
            mb_obs[batch_idx] = obs[e_idx]
            mb_obs_next[batch_idx] = get_obs_diff(obs[e_idx], o)
            mb_actions[batch_idx] = actions[e_idx]
            mb_rewards[batch_idx] = r

            total_reward[e_idx] += r
            total_steps[e_idx] += 1

            batch_idx = (batch_idx + 1) % BATCH_SIZE
            if batch_idx == 0:
                yield mb_obs, mb_obs_next, mb_actions, mb_rewards, done_rewards, done_steps
                done_rewards.clear()
                done_steps.clear()
            if done:
                o = e.reset()
                done_rewards.append(total_reward[e_idx])
                done_steps.append(total_steps[e_idx])
                total_reward[e_idx] = 0.0
                total_steps[e_idx] = 0
            obs[e_idx] = o

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--cuda", default=False, action="store_true", help="Enable cuda")
    parser.add_argument("-n", "--name", required=True, help="Name of the run")
    parser.add_argument("-m", "--model", required=True, help="File with model to load")
    args = parser.parse_args()
    device = torch.device("cuda" if args.cuda else "cpu")

    saves_path = os.path.join("saves", "02_env_" + args.name)
    os.makedirs(saves_path, exist_ok=True)

    envs = [common.make_env() for _ in range(NUM_ENVS)]
    writer = SummaryWriter(comment="-02_env_" + args.name)

    net = common.AtariA2C(envs[0].observation_space.shape, envs[0].action_space.n)
    net_em = i2a.EnvironmentModel(envs[0].observation_space.shape, envs[0].action_space.n).to(device)
    net.load_state_dict(torch.load(args.model, map_location=lambda storage, loc: storage))
    net = net.to(device)
    print(net_em)
    optimizer = optim.Adam(net_em.parameters(), lr=LEARNING_RATE)

    step_idx = 0
    best_loss = np.inf
    with ptan.common.utils.TBMeanTracker(writer, batch_size=100) as tb_tracker:
        for mb_obs, mb_obs_next, mb_actions, mb_rewards, done_rewards, done_steps in iterate_batches(envs, net, device):
            if len(done_rewards) > 0:
                m_reward = np.mean(done_rewards)
                m_steps = np.mean(done_steps)
                print("%d: done %d episodes, mean reward=%.2f, steps=%.2f" % (
                    step_idx, len(done_rewards), m_reward, m_steps))
                tb_tracker.track("total_reward", m_reward, step_idx)
                tb_tracker.track("total_steps", m_steps, step_idx)

            obs_v = torch.FloatTensor(mb_obs).to(device)
            obs_next_v = torch.FloatTensor(mb_obs_next).to(device)
            actions_t = torch.LongTensor(mb_actions.tolist()).to(device)
            rewards_v = torch.FloatTensor(mb_rewards).to(device)

            optimizer.zero_grad()
            out_obs_next_v, out_reward_v = net_em(obs_v.float()/255, actions_t)
            loss_obs_v = F.mse_loss(out_obs_next_v.squeeze(-1), obs_next_v)
            loss_rew_v = F.mse_loss(out_reward_v.squeeze(-1), rewards_v)
            loss_total_v = OBS_WEIGHT * loss_obs_v + REWARD_WEIGHT * loss_rew_v
            loss_total_v.backward()
            optimizer.step()
            tb_tracker.track("loss_em_obs", loss_obs_v, step_idx)
            tb_tracker.track("loss_em_reward", loss_rew_v, step_idx)
            tb_tracker.track("loss_em_total", loss_total_v, step_idx)

            loss = loss_total_v.data.cpu().numpy()
            if loss < best_loss:
                print("Best loss updated: %.4e -> %.4e" % (best_loss, loss))
                best_loss = loss
                fname = os.path.join(saves_path, "best_%.4e_%05d.dat" % (loss, step_idx))
                torch.save(net_em.state_dict(), fname)

            step_idx += 1
            if step_idx % SAVE_EVERY_BATCH == 0:
                fname = os.path.join(saves_path, "em_%05d_%.4e.dat" % (step_idx, loss))
                torch.save(net_em.state_dict(), fname)


# Distributed MNIST on grid based on TensorFlow MNIST example

from datetime import datetime
# A quick fix to run TF 1.X code in TF 2.X, we may want to properly migrate the Python script to TF 2.X API.
import tensorflow.compat.v1 as tf
tf.disable_v2_behavior()
from tensorflow.python.summary.writer.writer_cache import FileWriterCache as SummaryWriterCache
import math
import numpy
import json
import sys
from flink_ml_tensorflow.tensorflow_context import TFContext

def export_saved_model(sess, export_dir, tag_set, signatures):
    g = sess.graph
    g._unsafe_unfinalize()
    builder = tf.saved_model.builder.SavedModelBuilder(export_dir)

    signature_def_map = {}
    for key, sig in signatures.items():
        signature_def_map[key] = tf.saved_model.signature_def_utils.build_signature_def(
            inputs={name: tf.saved_model.utils.build_tensor_info(tensor) for name, tensor in
                    sig['inputs'].items()},
            outputs={name: tf.saved_model.utils.build_tensor_info(tensor) for name, tensor in
                     sig['outputs'].items()},
            method_name=sig['method_name'] if 'method_name' in sig else key)

        builder.add_meta_graph_and_variables(
            sess,
            tag_set.split(','),
            signature_def_map=signature_def_map,
            clear_devices=True)

        g.finalize()
        builder.save()

class ExportHook(tf.train.SessionRunHook):
    def __init__(self, export_dir, input_tensor, output_tensor):
        self.export_dir = export_dir
        self.input_tensor = input_tensor
        self.output_tensor = output_tensor

    def end(self, session):
        print("{} ======= Exporting to: {}".format(datetime.now().isoformat(), self.export_dir))
        signatures = {
            tf.saved_model.signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY: {
                'inputs': {'image': self.input_tensor},
                'outputs': {'prediction': self.output_tensor},
                'method_name': tf.saved_model.signature_constants.PREDICT_METHOD_NAME
            }
        }
        export_saved_model(session, self.export_dir, tf.saved_model.tag_constants.SERVING,
                           signatures)
        print("{} ======= Done exporting".format(datetime.now().isoformat()))

def decode(features):
    image = tf.decode_raw(features['image_raw'], tf.uint8)
    label = tf.one_hot(features['label'], 10, on_value=1)
    return image, label

def input_iter(context, batch_size):
    features = {'label': tf.FixedLenFeature([], tf.int64), 'image_raw': tf.FixedLenFeature([], tf.string)}
    dataset = context.flink_stream_dataset()
    dataset = dataset.map(lambda record: tf.parse_single_example(record, features=features))
    dataset = dataset.map(decode)
    dataset = dataset.batch(batch_size)
    iterator = tf.data.make_one_shot_iterator(dataset)
    return iterator

def map_fun(context):
    tf_context = TFContext(context)
    job_name = tf_context.get_role_name()
    task_index = tf_context.get_index()
    cluster_json = tf_context.get_tf_cluster()
    print (cluster_json)
    sys.stdout.flush()

    props = context.properties
    batch_size = int(props.get("batch_size"))
    checkpoint_dir = props.get("checkpoint_dir")
    export_dir = props.get("export_dir")

    # Parameters
    IMAGE_PIXELS = 28
    hidden_units = 128

    cluster = tf.train.ClusterSpec(cluster=cluster_json)
    server = tf.train.Server(cluster, job_name=job_name, task_index=task_index)

    def feed_dict(images, labels):
        xs = numpy.array(images)
        xs = xs.astype(numpy.float32)
        xs = xs / 255.0
        ys = numpy.array(labels)
        ys = ys.astype(numpy.uint8)
        return (xs, ys)

    if job_name == "ps":
        from time import sleep
        while True:
            sleep(1)
    elif job_name == "worker":

        # Assigns ops to the local worker by default.
        with tf.device(
                tf.train.replica_device_setter(worker_device="/job:worker/task:" + str(task_index), cluster=cluster)):

            # Placeholders or QueueRunner/Readers for input data
            x = tf.placeholder(tf.float32, [None, IMAGE_PIXELS * IMAGE_PIXELS], name="x")
            y_ = tf.placeholder(tf.float32, [None, 10], name="y_")

            # Variables of the hidden layer
            hid_w = tf.Variable(
                tf.truncated_normal([IMAGE_PIXELS * IMAGE_PIXELS, hidden_units], stddev=1.0 / IMAGE_PIXELS),
                name="hid_w")
            hid_b = tf.Variable(tf.zeros([hidden_units]), name="hid_b")
            hid_lin = tf.nn.xw_plus_b(x, hid_w, hid_b)
            hid = tf.nn.relu(hid_lin)

            # Variables of the softmax layer
            sm_w = tf.Variable(
                tf.truncated_normal([hidden_units, 10], stddev=1.0 / math.sqrt(hidden_units)),
                name="sm_w")
            sm_b = tf.Variable(tf.zeros([10]), name="sm_b")
            y = tf.nn.softmax(tf.nn.xw_plus_b(hid, sm_w, sm_b))

            global_step = tf.train.get_or_create_global_step()

            loss = -tf.reduce_sum(y_ * tf.log(tf.clip_by_value(y, 1e-10, 1.0)))

            train_op = tf.train.AdagradOptimizer(0.01).minimize(loss, global_step=global_step)

            # Test trained model
            label = tf.argmax(y_, 1, name="label")
            prediction = tf.argmax(y, 1, name="prediction")
            correct_prediction = tf.equal(prediction, label)

            accuracy = tf.reduce_mean(tf.cast(correct_prediction, tf.float32), name="accuracy")

            iter = input_iter(tf_context, batch_size)
            next_batch = iter.get_next()

            is_chief = (task_index == 0)
            sess_config = tf.ConfigProto(allow_soft_placement=True, log_device_placement=False,
                                         device_filters=["/job:ps", "/job:worker/task:%d" % task_index])

            # The MonitoredTrainingSession takes care of session initialization, restoring from
            #  a checkpoint, and closing when done or an error occurs
            mon_sess = tf.train.MonitoredTrainingSession(master=server.target, is_chief=is_chief,
                                                         checkpoint_dir=checkpoint_dir,
                                                         stop_grace_period_secs=10, max_wait_secs=300,
                                                         config=sess_config,
                                                         chief_only_hooks=[ExportHook(export_dir, x,
                                                                                prediction)])
            processed = 0
            while not mon_sess.should_stop():
                # Run a training step asynchronously
                # See `tf.train.SyncReplicasOptimizer` for additional details on how to
                # perform *synchronous* training.
                try:
                    images, labels = mon_sess.run(next_batch)
                    processed += images.shape[0]
                    # print mon_sess.run(next_batch)
                except tf.errors.OutOfRangeError:
                    break

                batch_xs, batch_ys = feed_dict(images, labels)
                feed = {x: batch_xs, y_: batch_ys}

                if len(batch_xs) > 0 and not mon_sess.should_stop():
                    _, step = mon_sess.run([train_op, global_step], feed_dict=feed)
                    if step % 100 == 0:
                        print("{0}, Task {1} step: {2} accuracy: {3}".format(
                            datetime.now().isoformat(), task_index, step,
                            mon_sess.run(accuracy, {x: batch_xs, y_: batch_ys})))
                        sys.stdout.flush()

            print(str(processed) + " records processed.")
            print("{0} stopping MonitoredTrainingSession".format(datetime.now().isoformat()))
            mon_sess.close()

    SummaryWriterCache.clear()

if __name__ == "__main__":
    map_fun(context)


# -*- coding: utf-8 -*-
# Original file: https://github.com/bolunwang/backdoor
# Original Author: Bolun Wang (bolunwang@cs.ucsb.edu)
# Original License: MIT
# Adpated to evaluate pytorch models

import os
import sys
import time

import numpy as np
from keras.preprocessing import image
import argparse
import seaborn as sns
import matplotlib.pyplot as plt
import pandas as pd
import pickle
from tensorboard.backend.event_processing import event_accumulator

parser = argparse.ArgumentParser(description='Neural Cleanse Testing')

# dataset and model
parser.add_argument('--dataset', type=str, default='zzz',
                    help='dataset to use')

parser.add_argument('--model-dir', type=str, default='./Neural_Cleanse/',
                    help='base dir')
parser.add_argument('--try-version', type=str, help='try version')
parser.add_argument('--save-version', type=int, default=4, help='save version')
parser.add_argument('--network-arch', type=str, help='network architecture')
parser.add_argument('--quantization', action='store_true')
parser.add_argument('--pruning', action='store_true')
parser.add_argument('--clean', action='store_true')
parser.add_argument('--tensorboard-path', type=str, default='../train_quantization_attack/tensorboard/')

args = parser.parse_args()

##############################
#        PARAMETERS          #
##############################
# NUM_CLASSES = 100  # total number of classes in the model

##############################
#      END PARAMETERS        #
##############################

def outlier_detection(l1_norm_list, idx_mapping):

    consistency_constant = 1.4826  # if normal distribution
    median = np.median(l1_norm_list)
    # print(l1_norm_list)
    mad = consistency_constant * np.median(np.abs(l1_norm_list - median))
    min_mad = np.abs(np.min(l1_norm_list) - median) / mad

    print('median: %f, MAD: %f' % (median, mad))
    print('anomaly index: %f' % min_mad)

    flag_list = []
    anomaly_index_list = []

    for y_label in idx_mapping:
        if l1_norm_list[idx_mapping[y_label]] > median:
            anomaly_index_list.append(0)
            continue
        else:
            anomaly_index_list.append(np.abs(l1_norm_list[idx_mapping[y_label]] - median) / mad)

        if np.abs(l1_norm_list[idx_mapping[y_label]] - median) / mad > 2:
            flag_list.append((y_label, l1_norm_list[idx_mapping[y_label]]))

    if len(flag_list) > 0:
        flag_list = sorted(flag_list, key=lambda x: x[1])

    print('flagged label list: %s' %
          ', '.join(['%d: %2f' % (y_label, l_norm)
                     for y_label, l_norm in flag_list]))

    return anomaly_index_list

def analyze_pattern_norm_dist():

    mask_flatten = []
    idx_mapping = {}

    mask_flatten_full = []

    l1_average_norm_list = []

    for y_label in range(NUM_CLASSES):
        tmp_mask = None
        tmp_l1_norm_list = []
        for round_ in range(3):
            mask_filename = IMG_FILENAME_TEMPLATE % ('mask', y_label, round_)
            # print(mask_filename)
            if os.path.isfile('%s/%s' % (RESULT_DIR, mask_filename)):
                img = image.load_img(
                    '%s/%s' % (RESULT_DIR, mask_filename),
                    color_mode='grayscale',
                    target_size=INPUT_SHAPE)
                mask = image.img_to_array(img)
                mask /= 255
                mask = mask[:, :, 0]
                mask_flatten_full.append(mask.flatten())
                tmp_l1_norm_list.append(np.sum(np.abs(mask.flatten())))

        l1_average_norm_list.append(min(tmp_l1_norm_list))
        idx_mapping[y_label] = len(l1_average_norm_list) - 1

    l1_norm_list = l1_average_norm_list

    l1_norm_list_full = [np.sum(np.abs(m)) for m in mask_flatten_full]

    print('%d labels found' % len(l1_norm_list))

    return outlier_detection(l1_norm_list, idx_mapping)

if __name__ == '__main__':
    print('%s start' % sys.argv[0])
    if args.dataset == 'cifar10':
        NUM_CLASSES = 10

    elif args.dataset == 'cifar100':
        NUM_CLASSES = 100

    elif args.dataset == 'gtsrb':
        NUM_CLASSES = 43

    save_version = args.save_version

    network_arch = args.network_arch
    result_list = []
    if args.clean:
        ckpt_file_num = 10
    else:
        ckpt_file_num = 5
    for target_label in range(ckpt_file_num):
        if not args.clean:
            target_label = 2 * target_label + 1
        if args.dataset == 'cifar10':
            dataset_name = 'CIFAR10'
        elif args.dataset == 'gtsrb':
            dataset_name = 'GTSRB'

        if args.pruning:
            ckpt_name_template = "%s_%s_label_%d_%s_v%d/best.pth"
            RESULT_DIR = ckpt_name_template % (network_arch, dataset_name.upper(), target_label, args.try_version, save_version)
        elif args.quantization:
            ckpt_folder = "quantization_%s_%s_100_c_1_Mlayers4_v%s_label%d"  % (network_arch, dataset_name.upper(), args.try_version, target_label)
            RESULT_DIR = ckpt_folder

        elif args.clean:
            RESULT_DIR = '%s_%s_clean_v%d/best.pth' % (network_arch, dataset_name.lower(), target_label)
        else:
            raise ValueError('Unknown setting')

        RESULT_DIR = RESULT_DIR.replace('/', '_')
        RESULT_DIR = args.model_dir + RESULT_DIR
        IMG_FILENAME_TEMPLATE = args.dataset + '_visualize_%s_label_%d_v%d.png'  # image filename template for visualization results

        # input size
        IMG_ROWS = 32
        IMG_COLS = 32
        IMG_COLOR = 3
        INPUT_SHAPE = (IMG_ROWS, IMG_COLS, IMG_COLOR)

        start_time = time.time()
        print(RESULT_DIR)
        result = analyze_pattern_norm_dist()
        elapsed_time = time.time() - start_time
        print('elapsed time %.2f s' % elapsed_time)
        result_list.append(result)

    print(result_list)

    if not os.path.exists('./neural_cleanse_anomaly_index/'):
        os.makedirs('./neural_cleanse_anomaly_index/')
    anomaly_index_save_path = './neural_cleanse_anomaly_index/' + '%s_%s_%s' % (network_arch, dataset_name.lower(), args.try_version)

    with open(anomaly_index_save_path, 'wb') as f:
        pickle.dump(result_list, f)

    print('save at %s'% anomaly_index_save_path)