text:
#!/usr/local/bin/python
# -*- coding: utf-8 -*-
import re
import numpy as np
import os
import glob

# from HDMap_pb2 import HDMap
# from proto.SourceInfo_pb2 import GPSInfo

class element(object):
    # uint8 id
    # uint8 color
    # uint8 element_type
    # sensor_msgs/PointCloud2 pointclouds
    # uint8[] direction
    def __init__(self):
        self.id = 0
        self.color = 0
        self.element_type = 0
        self.direction = []
    def copy_ele(self, ele):
        self.id = ele.id
        self.color = ele.color
        self.element_type = ele.element_type
        self.direction = ele.direction

class elements(object):
    # uint8 type
    # uint8 motor_type
    # element[] elements
    def __init__(self):
        self.type = 0
        self.motor_type = 0
        self.elements = element

    def copy_ele(self, eles):
        self.type = eles.type
        self.motor_type = eles.motor_type
        self.elements = eles.element

class Divider(object):
    def __init__(self, id, color, type, geometry):
        self.id = id
        # 颜色，0：未知；1：白色；2：黄色
        self.color = color
        # 类型 0: 未知； 1：单虚线；2：单实线；3：双实线；4：左实右虚；5：左虚右实,；6: 不可通行减速线；7：可通行减速线
        self.type = type
        self.geometry = geometry

    def getPoints(self, normal = True, stacked=False):
        elestr = self.geometry
        tmp = re.findall(r"\d+\.?\d*", elestr)
        point_lenth = 3
        # [lat, lon, 0]
        temp_point =np.array([tmp[i:i + point_lenth] for i in range(0, len(tmp), point_lenth)], dtype=float)
        # print("temp_point", temp_point)
        if normal == True:
            temp_point = normalize_for_rviz(temp_point)
        if stacked == False:
            return temp_point # shape = (n, 3)
        x = temp_point[:, 0].reshape(-1)
        y = temp_point[:, 1].reshape(-1)
        z = temp_point[:, 2].reshape(-1)

        points = np.stack((x, y, z))

        return points #shape = (3, n)

class Trafficlights(element):
    def __init__(self, ele):
        self.id = ele.id
        self.geometry = ele.geometry
        self.motor_type = ele.motor_type
        self.lights = ele.lights
        self.geomtry_record = []
        self.location = []

    def getPoints(self, normal = True):
        elestr = self.geometry
        tmp = re.findall(r"\d+\.?\d*", elestr)
        self.location.append(tmp)
        for dirs in self.lights:
            #  方向，1：竖排；2：横排; 进入ros后1变为49,2变为50
            # self.direction += str(dirs.direction)
            # read point
            _elestr = dirs.geometry
            tmp = re.findall(r"\d+\.?\d*", _elestr)
            _id = dirs.id
            self.geomtry_record.append((_id, dirs.direction, tmp))

        if normal == True:
            # print(np.shape(self.location))
            self.location = normalize_for_rviz(self.location)

        return self.location
class Landmarks(element):
    def __init__(self, ele):
        self.id = ele.id
        self.type = ele.type
        self.geometry = ele.geometry

    def getPoints(self, normal = True):
        elestr = self.geometry
        tmp = re.findall(r"\d+\.?\d*", elestr)
        step = 3
        temp_point = np.array([tmp[i:i + step] for i in range(0, len(tmp), step)], dtype=float)

        if normal == True:
            temp_point = normalize_for_rviz(temp_point)

        return temp_point

def read_gps_all(path):
    points = []
    headings = []
    count = 0

    for f in glob.glob(path+"/gps_store/*.txt"): # glob.glob(path): #
        #if count ==2:
        #    break
        count+=1
        # print(count)
        #if count ==2:
        #    break
        tmp = np.loadtxt(f, delimiter=',')
        points_line = []
        headings_line = []
        for ele in tmp:
            points_line.append(ele[:3])
            headings_line.append(ele[3])
        #    print('points_line',points_line)
        #    print('headings_line', headings_line)
        # print('ele',np.shape(ele))
        # print('np.shape(headings_line)', np.shape(headings_line))
        points.append(np.array(points_line))
        headings.append(np.array(headings_line))

        # print('np.shape(headings)', np.shape(headings))

    return np.array(points), np.array(headings)

def read_from_pb(path = "package://rviz_tools_py/data", write2txt = False):
    message = HDMap()
    with open(path+'/hdmap_deecamp.pb', 'rb') as fb:
        pb_content = fb.read() # 如果文件很大，则要分批读取
        # pb_content 是二进制的pb数据，如果是文件，则需要用 open 方法读取数据
        message.ParseFromString(pb_content)
    # print(message)

    if write2txt==True:
        with open(path+'/hdmap_message.txt', 'w') as fb:
            print >> fb, message
    return message

class Divider2(object):
    def __init__(self, id, color, type, geometry):
        self.id = id
        # 颜色，0：未知；1：白色；2：黄色
        self.color = color
        # 类型 0: 未知； 1：单虚线；2：单实线；3：双实线；4：左实右虚；5：左虚右实,；6: 不可通行减速线；7：可通行减速线
        self.type = type
        self.geometry = geometry

    # def getPoints(self, normal = True, stacked=False):
    #     elestr = self.geometry
    #     tmp = re.findall(r"\d+\.?\d*", elestr)
    #     point_lenth = 3
    #     # [lat, lon, 0]
    #     temp_point =np.array([tmp[i:i + point_lenth] for i in range(0, len(tmp), point_lenth)], dtype=float)
    #     # print("temp_point", temp_point)
    #     if normal == True:
    #         temp_point = normalize_for_rviz(temp_point)
    #     if stacked == False:
    #         return temp_point # shape = (n, 3)
    #     x = temp_point[:, 0].reshape(-1)
    #     y = temp_point[:, 1].reshape(-1)
    #     z = temp_point[:, 2].reshape(-1)

    #     points = np.stack((x, y, z))

    #     return points #shape = (3, n)

def read_divider_txt(file_path = "/home/lfg/my_work/deecamp/lfg/HDMapProject/rviz_ws/src/rviz_tools_py/data/divider_hdmap.txt"):
    id_list = []
    type_list = []
    color_list = []
    geo_list = []

    with open(file_path, 'rt') as f:
        lines = [line.strip('\n').strip(',') for line in f]
        # print(np.shape(lines))
        line_length = 0
        for line in lines:
            line = re.split(',', line)
            # print(float(_line[0]))

            try:
                f_line = [float(ele) for ele in line]
            except ValueError:
                print(line)
                print("error!")
                break
            #print(line)
            #print(f_line)
            id_list.append(f_line[0])
            type_list.append(f_line[1])
            color_list.append(f_line[2])

            step = 3
            _geo = np.array([f_line[i:i+step] for i in range(line_length+3,len(f_line), step)], dtype = float)

            # fix bug of txt
            line_length = len(f_line)-3
            # print('len(f_line)', len(_geo))
            # print("lenth", line_length)

            # print(_geo)
            geo_list.append(_geo)

    divider_list = []
    for i in range(len(id_list)):
        divider_list.append(Divider2(id_list[i], color_list[i], type_list[i], np.array(geo_list[i])))

    return np.array(geo_list)
    # return divider_list

class Gps(object):
    def __init__(self, scene_id, device_id, pts):
        self.scene_id = scene_id
        self.device_id = device_id
        self.pts = Pts(pts.id, pts.gpstime, pts.heading, pts.speed, pts.geometry)

class Pts(object):
    def __init__(self, id, gpstime, heading, speed, geometry):
        self.id = float(id)
        self.gpstime = float(gpstime)
        self.heading = float(heading)
        self.speed = float(speed)
        self.geometry = self.getGeometry(geometry)

    def getGeometry(self, geometry):

        tmp = re.findall(r"\d+\.?\d*", geometry)
        t_geo = [float(i) for i in tmp]

        return t_geo

# from proto.SourceInfo_pb2 import GPSInfo
def read_gps_txt(file_path='"package://rviz_tools_py/data/20190123112752_7ac6ab9d61d94314188426910d324c39_4deecamp_gps.pb'):
    message = GPSInfo()
    with open(file_path, 'rt') as f:
        pb_content = f.read() # 如果文件很大，则要分批读取
        # pb_content 是二进制的pb数据，如果是文件，则需要用 open 方法读取数据
        message.ParseFromString(pb_content)

    gps = Gps(message.scene_id, message.device_id, message.pts)
    return gps

def read_from_pb_debug(path = '/home/lfg/my_work/deecamp/lfg/HDMapProject/rviz_ws/src/rviz_tools_py/data/hdmap_deecamp.pb', write2txt = False):
    message = HDMap()
    with open(path, 'rb') as fb:
        pb_content = fb.read() # 如果文件很大，则要分批读取
        # pb_content 是二进制的pb数据，如果是文件，则需要用 open 方法读取数据
        message.ParseFromString(pb_content)
    # print(message)

    if write2txt==True:
        with open(path+'/hdmap_message.txt', 'w') as fb:
            print >> fb, message
    return message

def normalize_for_rviz(temp_point):
    # lat_max = np.max(temp_point[:, 0])
    # lon_max = np.max(temp_point[:, 1])
    # lat_min = np.min(temp_point[:, 0])
    # lon_min = np.min(temp_point[:, 1])

    # print("lat_min, lon_min", lat_min, lon_min)
    # print("temp_point[:, 0][3]", temp_point[:, 0][0])
    # print("temp_point[:, 1][3]", temp_point[:, 1][0])
    lat_min = 22.68037672
    lon_min = 114.36243453
    if type(temp_point) == list:
        temp_point = np.array(temp_point, dtype = float)
    temp_point[:, 0] = 100000*(temp_point[:, 0]-lat_min)
    temp_point[:, 1] = 100000*(temp_point[:, 1]-lon_min)
    # print("After temp_point[:, 0][3]", temp_point[:, 0][:])
    # print("After temp_point[:, 1][3]", temp_point[:, 1])

    return temp_point

def read_divider(message):
    points_list = []
    for ele in message.dividers:
        this_ele = Divider(ele.id, ele.color, ele.type, ele.geometry)

        points= this_ele.getPoints(normal = True, stacked = False)

        # print(points.shape)
        points_list.append(points)

    return points_list

def read_trafficlight(message):
    points_list = []
    for ele in message.tafficlights:
        trafficlight = Trafficlights(ele)

        points = trafficlight.getPoints()
        points_list.append(points)

    return points_list

def read_lanemarkings(message):
    points_list = []
    for ele in message.lanemarkings:
        lanemarkings = Landmarks(ele)

        points = lanemarkings.getPoints()
        points_list.append(points)

    return points_list

gps_test_path = '/home/apple/catkin_ws/src/rviz_tools_py/data/GPStrafficLine.txt'
test_bug_txt = '/home/apple/catkin_ws/test_bug.txt'
np_data_path = '/home/apple/catkin_ws/src/rviz_tools_py/data/gps_pre_frame.npy'

def read_points_per_frame(path='/home/apple/catkin_ws/src/rviz_tools_py/data/GPStrafficLine.txt'):

    tmp = np.loadtxt(gps_test_path, delimiter=' ')
    frame = []
    points_line = []
    points_line_i = []
    before = 0
    for ele in tmp:

        if all(ele[:2] == before):
            points_line_i.append(ele[2:4])
        elif before !=0:
            frame.append(ele[:2])
            before  = ele[:2]
            points_line.append(points_line_i)
            points_line_i = []
            points_line_i.append(ele[2:4])
        elif before ==0:
            frame.append(ele[:2])
            points_line_i.append(ele[2:4])

    points_line.append(points_line_i)

    return frame, points_line

# if __name__ == "__main__":
#     message = read_from_pb_debug()
#     # res = read_divider(message)
#     res  = read_lanemarkings(message)
#     print(res)

    # print("result:   ", res[0])

max_stars_repo_name: LiuFG/UpdatingHDmapByMonoCamera

max_issues_repo_name: DeeCamp-Demo/HDMapProject

max_forks_repo_name: DeeCamp-Demo/HDMapProject


text:
import cv2
import concurrent.futures
import os
import numpy as np
import json
import copy
import random
import pickle

def get_root_path():
    return os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))

def sliding_crop_canny_imgs(img_dir,
                            save_crop_dir,
                            sliding_win_xsize=1650,
                            sliding_win_ysize=1650,
                            overlap=200):

    def crop(img_dir, img_name, save_crop_dir, sliding_win_xsize,
             sliding_win_ysize, overlap):
        full_path = os.path.join(img_dir, img_name)
        image = cv2.imread(full_path)
        resized_image = cv2.resize(
            image, (int(0.1 * image.shape[1]), int(0.1 * image.shape[0])))
        imgray = cv2.cvtColor(resized_image, cv2.COLOR_BGR2GRAY)
        ret, bin_pic = cv2.threshold(imgray, 0, 255,
                                     cv2.THRESH_BINARY + cv2.THRESH_OTSU)
        median = cv2.medianBlur(bin_pic, 5)
        cannyPic = cv2.Canny(median, 10, 200)

        contours, hierarchy = cv2.findContours(cannyPic, cv2.RETR_CCOMP,
                                               cv2.CHAIN_APPROX_SIMPLE)

        maxArea = 0
        for i in range(len(contours)):
            if cv2.contourArea(contours[i]) > cv2.contourArea(
                    contours[maxArea]):
                maxArea = i

        boader = contours[maxArea]
        xmin, xmax, ymin, ymax = np.min(boader[:, :, 0]), np.max(
            boader[:, :, 0]), np.min(boader[:, :, 1]), np.max(boader[:, :, 1])

        xmin = xmin * 10 - 150
        xmax = xmax * 10 + 150
        ymin = ymin * 10 - 150
        ymax = ymax * 10 + 150

        if xmax - xmin > 1500 and ymax - ymin > 1500 and xmin >= 0 and xmax <= image.shape[
                1] and ymin >= 0 and ymax <= image.shape[0]:
            croped_img = image[ymin:ymax + 1, xmin:xmax + 1]
        else:
            # if the canny croped is invalid, change to center crop
            center = [image.shape[1] // 2, image.shape[0] // 2]
            if 'CAM3' in img_name:
                center[1] = int(image.shape[0] * 0.55)
            length = int(image.shape[1] * 0.36)
            croped_img = image[center[1] - length:center[1] + length + 1,
                               center[0] - length:center[0] + length + 1, :]
            xmin, xmax, ymin, ymax = center[0] - length, center[
                0] + length, center[1] - length, center[1] + length

        height = ymax - ymin + 1
        width = xmax - xmin + 1
        x_slides = []
        y_slides = []
        x_slide = 0
        y_slide = 0
        while True:
            if x_slide >= width - sliding_win_xsize:
                x_slide = width - sliding_win_xsize
                x_slides.append(x_slide)
                break
            else:
                x_slides.append(x_slide)
                x_slide += (sliding_win_xsize - overlap)

        while True:
            if y_slide >= height - sliding_win_ysize:
                y_slide = height - sliding_win_ysize
                y_slides.append(y_slide)
                break
            else:
                y_slides.append(y_slide)
                y_slide += (sliding_win_ysize - overlap)

        slides_pos = [(x + xmin, y + ymin) for x in x_slides for y in y_slides]
        for slide_pos in slides_pos:
            x, y = slide_pos
            croped_img = image[y:y + sliding_win_ysize,
                               x:x + sliding_win_xsize]
            croped_img_name = '{}_{}_{}_{}_{}.jpg'.format(
                img_name.split('.')[0], x, y, x + sliding_win_xsize - 1,
                y + sliding_win_ysize - 1)
            cv2.imwrite(
                os.path.join(save_crop_dir, croped_img_name), croped_img)

        print('Finish {}'.format(img_name))

    img_names = os.listdir(img_dir)
    if not os.path.exists(save_crop_dir):
        os.mkdir(save_crop_dir)

    with concurrent.futures.ThreadPoolExecutor(max_workers=40) as executor:
        for img_name in img_names:
            executor.submit(crop, img_dir, img_name, save_crop_dir,
                            sliding_win_xsize, sliding_win_ysize, overlap)

def generate_canny_slide_annotations_file(anno_file, img_dir, save_anno_folder,
                                          train_anno_json_name,
                                          bk_anno_json_name,
                                          all_anno_json_name):
    '''
    params:
        anno_file: original annotation json file path
        img_dir: the processed canny slide img patches folder
        save_anno_folder: the folder for saving the generated annotations
        train_anno_json_name: name for generated train annotations (with defects, positive samples)
        bk_anno_json_name: name for all background annotations
        all_anno_json_name: name for all the samples(all positive and negative) annotations
    '''

    with open(anno_file, 'r') as f:
        annotations = json.load(f)

    name_anno_map = {}
    for anno in annotations:
        name = anno['name']
        if name not in name_anno_map.keys():
            name_anno_map[name] = [anno]
        else:
            name_anno_map[name].append(anno)

    all_annotations = []
    new_annotations = []
    background_annotations = []

    img_names = os.listdir(img_dir)

    for img_name in img_names:
        has_box = False
        splits = (img_name.split('.')[0]).split('_')
        origin_img_name = '{}_{}_{}_{}.jpg'.format(splits[0], splits[1],
                                                   splits[2], splits[3])
        xmin, ymin, xmax, ymax = int(splits[-4]), int(splits[-3]), int(
            splits[-2]), int(splits[-1])
        annos = name_anno_map[origin_img_name]
        image_width = xmax - xmin + 1
        image_height = ymax - ymin + 1

        for ann in annos:
            anno = copy.deepcopy(ann)

            anno['image_width'] = image_width
            anno['image_height'] = image_height
            anno['name'] = img_name
            left_up_position = [xmin, ymin]
            right_down_position = [xmax, ymax]
            box = anno['bbox']

            if box[0] >= left_up_position[0] and box[0] <= right_down_position[
                    0] and box[1] >= left_up_position[1] and box[
                        1] <= right_down_position[1] and box[
                            2] >= left_up_position[0] and box[
                                2] <= right_down_position[0] and box[
                                    3] >= left_up_position[1] and box[
                                        3] <= right_down_position[1]:
                has_box = True
                anno['bbox'][0] = round(
                    max(0, anno['bbox'][0] - left_up_position[0]), 2)  # x1
                anno['bbox'][1] = round(
                    max(0, anno['bbox'][1] - left_up_position[1]), 2)  # y1
                anno['bbox'][2] = round(
                    max(0, anno['bbox'][2] - left_up_position[0]), 2)  # x2
                anno['bbox'][3] = round(
                    max(0, anno['bbox'][3] - left_up_position[1]), 2)  # y2
                new_annotations.append(anno)

        if not has_box:
            anno = dict(
                name=img_name,
                image_height=image_height,
                image_width=image_width,
                category=0,
                bbox=[])
            background_annotations.append(anno)

    select_indices = random.sample(
        range(len(background_annotations)), len(new_annotations))
    bks = [background_annotations[idx] for idx in select_indices]
    all_annotations = bks + new_annotations
    with open(os.path.join(save_anno_folder, train_anno_json_name), 'w') as f:
        json.dump(new_annotations, f, indent=4, separators=(',', ': '))
    with open(os.path.join(save_anno_folder, bk_anno_json_name), 'w') as f:
        json.dump(background_annotations, f, indent=4, separators=(',', ': '))
    with open(os.path.join(save_anno_folder, all_anno_json_name), 'w') as f:
        json.dump(all_annotations, f, indent=4, separators=(',', ': '))

def test_predictions(predict_json_file, image_dir, result_img_save_dir):
    '''
    Descriptions:
        function to test the generated submit json file
    '''
    with open(predict_json_file, 'r') as f:
        context = json.load(f)

    image_names = os.listdir(image_dir)
    # select_indices = random.sample(range(len(image_names)), 10)
    select_indices = range(10)
    select_img_names = [image_names[i] for i in select_indices]

    if not os.path.exists(result_img_save_dir):
        os.mkdir(result_img_save_dir)

    for index, image_name in enumerate(select_img_names):
        image = cv2.imread(os.path.join(image_dir, image_name))
        for con in context:
            if con['name'] == image_name:
                cv2.rectangle(image,
                              (int(con['bbox'][0]), int(con['bbox'][1])),
                              (int(con['bbox'][2]), int(con['bbox'][3])),
                              (0, 0, 255))
                cv2.putText(
                    image,
                    str(con['category']) + '_' + str(round(con['score'], 2)),
                    (int(con['bbox'][0]), int(con['bbox'][1]) - 10),
                    cv2.FONT_HERSHEY_COMPLEX, 1, (255, 0, 0), 1)
        cv2.imwrite(
            os.path.join(result_img_save_dir, 'test{}.jpg'.format(index)),
            image)

def convert_tile_annotations_toinfos(ann_file, val_num_rate, infos_save_dir,
                                     train_info_name, val_info_name,
                                     all_info_name):
    '''
    Description:
        function to convert from the anno json to the mmdetection dataset annotation infos,
        saved in pickle format
    '''
    with open(ann_file, 'r') as f:
        text = json.load(f)

    if not os.path.exists(infos_save_dir):
        os.mkdir(infos_save_dir)

    total_annotations = {}
    total_infos = []

    for annotation in text:
        if annotation == "info" or annotation == "licenses":
            continue
#         import pdb;pdb.set_trace()
        if annotation['file_name'] not in total_annotations.keys():
            total_annotations[annotation['file_name']] = [
                annotation['height'], annotation['width'],
                [annotation['category_id']], [annotation['bbox']]
            ]
        else:
            total_annotations[annotation['file_name']][2].append(
                annotation['category_id'])
            total_annotations[annotation['file_name']][3].append(annotation['bbox'])

    for file_name in total_annotations.keys():
        file_info = total_annotations[file_name]
        bboxes = file_info[-1]
        labels = file_info[-2]
        if len(bboxes[0]) == 0:
            bboxes = np.zeros((0, 4), dtype=np.float32)
            labels = np.array([], dtype=np.int64)
        else:
            bboxes = np.array(bboxes).astype(np.float32)
            labels = np.array(labels).astype(np.int64)
        total_infos.append(
            dict(
                filename=file_name,
                height=file_info[0],
                width=file_info[1],
                ann=dict(bboxes=bboxes, labels=labels)))

    val_num = int(len(total_infos) * val_num_rate)
    val_infos_indices = random.sample(range(len(total_infos)), val_num)
    train_info_indices = list(
        set(range(len(total_infos))).difference(set(val_infos_indices)))

    train_infos = [total_infos[idx] for idx in train_info_indices]
    val_infos = [total_infos[idx] for idx in val_infos_indices]

    with open(os.path.join(infos_save_dir, train_info_name), 'wb') as f:
        pickle.dump(train_infos, f)
    with open(os.path.join(infos_save_dir, val_info_name), 'wb') as f:
        pickle.dump(val_infos, f)
    with open(os.path.join(infos_save_dir, all_info_name), 'wb') as f:
        pickle.dump(total_infos, f)

    test_recover_infos(
        val_info_file=os.path.join(infos_save_dir, val_info_name),
        train_info_file=os.path.join(infos_save_dir, train_info_name),
        all_info_file=os.path.join(infos_save_dir, all_info_name))

def test_recover_infos(val_info_file, train_info_file, all_info_file):
    with open(val_info_file, 'rb') as f:
        val_infos = pickle.load(f)
    with open(train_info_file, 'rb') as f:
        train_infos = pickle.load(f)
    with open(all_info_file, 'rb') as f:
        all_infos = pickle.load(f)
    assert (type(val_infos[0]) == dict)
    print(val_infos[0], len(val_infos))
    print(train_infos[0], len(train_infos))
    print(len(all_infos))

def check_annotations(anno_json_path):
    with open(anno_json_path, 'r') as f:
        context = json.load(f)

    for anno in context:
        img_name = anno['name']
        box = anno['bbox']
        height = anno['image_height']
        width = anno['image_width']
        if len(box) > 0:
            if box[0] < 0 or box[0] >= width or box[2] < 0 or box[
                    2] >= width or box[1] < 0 or box[1] >= height or box[
                        3] < 0 or box[3] >= height or box[0] >= box[2] or box[
                            1] >= box[3]:
                print(anno)

ext: py

max_stars_repo_name: NEUdeep/TileDetection

max_issues_repo_name: NEUdeep/TileDetection

max_forks_repo_name: NEUdeep/TileDetection

reason: import numpy


text:
import sys
sys.path += ['../utils']
import csv
from tqdm import tqdm
import collections
import gzip
import pickle
import numpy as np
import faiss
import os
import pytrec_eval
import json
from msmarco_eval import quality_checks_qids, compute_metrics, load_reference

  # location for dumpped query and passage/document embeddings which is output_dir
#checkpoint_path ='/home/dihe/cudnn_file/recommender_shuqi/MIND_data/raw_data/exp_12_02_04/ann_data/'
# checkpoint =  150000 # embedding from which checkpoint(ie: 200000)
# data_type = 0 # 0 for document, 1 for passage
# test_set = 1 # 0 for dev_set, 1 for eval_set
# raw_data_dir = '/home/dihe/cudnn_file/recommender_shuqi/MIND_data/raw_data/'
# processed_data_dir = '/home/dihe/cudnn_file/recommender_shuqi/MIND_data/raw_data/ann_data_roberta-base-fast-doc_512'

# checkpoint_path ='/home/dihe/Projects/data/raw_data/exp_12_02_04/ann_data/'
# raw_data_dir = '/home/dihe/Projects/data/raw_data/'
# processed_data_dir = '/home/dihe/Projects/data/raw_data/ann_data_roberta-base-fast-doc_512'

# checkpoint =  0
# data_type = 0
# test_set = 1
# checkpoint_path ='/home/dihe/Projects/data/raw_data/test_roberta_decode_doc/ann_data/'
# raw_data_dir = '/home/dihe/Projects/data/raw_data/'
# processed_data_dir = '/home/dihe/Projects/data/raw_data/ann_data_roberta-base-fast-doc_512'

#--------------------------------------------------------------------------------------
# checkpoint =  0
# data_type = 0
# test_set = 0
# checkpoint_path ='/home/dihe/cudnn_file/recommender_shuqi/MIND_data/raw_data/exp_12_19_01/ann_data2/'
# raw_data_dir = '/home/dihe/Projects/data/raw_data/'
# processed_data_dir = '/home/dihe/Projects/data/raw_data/ann_data_roberta-base-fast-docdev_512'

# checkpoint =  0
# data_type = 0
# test_set =0
# checkpoint_path ='/home/dihe/cudnn_file/recommender_shuqi/MIND_data/raw_data/exp_12_23_02/ann_data400000/'
# raw_data_dir = '/home/dihe/Projects/data/raw_data/'
# processed_data_dir = '/home/dihe/Projects/data/raw_data/ann_data_roberta-base-fast-docdev_512'

# checkpoint =  0
# data_type = 0
# test_set = 0
# checkpoint_path ='/home/dihe/cudnn_file/recommender_shuqi/MIND_data/raw_data/exp_12_23_02/ann_data4/'
# raw_data_dir = '/home/dihe/Projects/data/raw_data/'
# processed_data_dir = '/home/dihe/Projects/data/raw_data/ann_data_roberta-base-fast-docdev_512'

# checkpoint =  0
# data_type = 0
# test_set = 1
# checkpoint_path ='/home/dihe/cudnn_file/recommender_shuqi/MIND_data/raw_data/exp_01_05_09/ann_data820000/'
# raw_data_dir = '/home/dihe/Projects/data/raw_data/'
# processed_data_dir = '/home/dihe/Projects/data/raw_data/ann_data_roberta-base-fast-docdev2_512'
# # processed_data_dir2 = '/home/dihe/Projects/data/raw_data/ann_data_roberta-base-fast-docdev2_512'
# # checkpoint_path2 ='/home/dihe/cudnn_file/recommender_shuqi/MIND_data/raw_data/exp_01_05_09/ann_data820000/'

# processed_data_dir2 = '/home/dihe/Projects/data/raw_data/ann_data_roberta-base-fast-doceval_512'
# checkpoint_path2 ='/home/dihe/Projects/data/raw_data/eval_exp_01_05_09_820000/ann_data/'
# query_emb_num=4

checkpoint =  0
data_type = 0
test_set = 1
# checkpoint_path ='/home/dihe/cudnn_file/recommender_shuqi/MIND_data/raw_data/exp_01_05_09/ann_data820000/'
# checkpoint_path ='/home/dihe/Projects/data/raw_data/exp_01_05_09/ann_data910000/'
# raw_data_dir = '/home/dihe/Projects/data/raw_data/'
# processed_data_dir = '/home/dihe/Projects/data/raw_data/ann_data_roberta-base-fast-docdev_512'
# processed_data_dir2 = '/home/dihe/Projects/data/raw_data/ann_data_roberta-base-fast-docdev2_512'
# checkpoint_path2 ='/home/dihe/cudnn_file/recommender_shuqi/MIND_data/raw_data/exp_01_05_09/ann_data820000/'

#training mrr
# checkpoint_path ='/home/dihe/Projects/data/raw_data/exp_01_05_09/ann_data910000/'
# raw_data_dir = '/home/dihe/Projects/data/raw_data/'
# processed_data_dir = '/home/dihe/Projects/data/raw_data/ann_data_roberta-base-fast-docdev_512'

# processed_data_dir2 = '/home/dihe/Projects/data/raw_data/ann_data_roberta-base-fast-trainqueryeval_512'
# checkpoint_path2 ='/home/dihe/Projects/data/raw_data/eval_exp_01_05_09_910000/ann_data/'
# query_emb_num=4

# checkpoint_path ='/home/dihe/Projects/data/raw_data/exp_01_05_09/ann_data910000/'
# raw_data_dir = '/home/dihe/Projects/data/raw_data/'
# processed_data_dir = '/home/dihe/Projects/data/raw_data/ann_data_roberta-base-fast-docdev_512'

# processed_data_dir2 = '/home/dihe/Projects/data/raw_data/ann_data_roberta-base-fast-trainqueryeval2_512'
# checkpoint_path2 ='/home/dihe/Projects/data/raw_data/eval_exp_01_05_09_910000/ann_data_sample20q/'
# query_emb_num=4

# checkpoint_path ='/home/dihe/Projects/data/raw_data/exp_21_05_21_01_check/ann_data10000/'
# raw_data_dir = '/home/dihe/Projects/data/raw_data/'
# processed_data_dir = '/home/dihe/Projects/data/raw_data/ann_data_roberta-base-fast-docdev_512'

# processed_data_dir2 = '/home/dihe/Projects/data/raw_data/ann_data_roberta-base-fast-trainqueryeval_512'
# checkpoint_path2 ='/home/dihe/Projects/data/raw_data/exp_21_05_21_01_check_10000/ann_data_sampleq/'
# query_emb_num=4

# checkpoint_path ='/home/dihe/cudnn_file/recommender_shuqi/MIND_data/raw_data/exp_21_05_21_01/check/ann_data10000/'
# raw_data_dir = '/home/dihe/Projects/data/raw_data/'
# processed_data_dir = '/home/dihe/Projects/data/raw_data/ann_data_roberta-base-fast-docdev_512'

# processed_data_dir2 = '/home/dihe/Projects/data/raw_data/ann_data_roberta-base-fast-trainqueryeval2_512'
# checkpoint_path2 ='/home/dihe/Projects/data/raw_data/exp_21_05_21_01_check_10000/ann_data_sample20q/'
# query_emb_num=4

# checkpoint_path ='/home/dihe/Projects/data/raw_data/exp_01_05_09/ann_data900000/'
# raw_data_dir = '/home/dihe/Projects/data/raw_data/'
# processed_data_dir = '/home/dihe/Projects/data/raw_data/ann_data_roberta-base-fast-docdev_512'

# processed_data_dir2 = '/home/dihe/Projects/data/raw_data/ann_data_roberta-base-fast-trainqueryeval2_512'
# checkpoint_path2 ='/home/dihe/Projects/data/raw_data/eval_exp_01_05_09_900000/ann_data_sample20q/'
# query_emb_num=4

# checkpoint_path ='/home/dihe/cudnn_file/recommender_shuqi/MIND_data/raw_data/exp_21_04_14_01/check/ann_data1000000/'
# raw_data_dir = '/home/dihe/Projects/data/raw_data/'
# processed_data_dir = '/home/dihe/Projects/data/raw_data/ann_data_roberta-base-fast-docdev_512'

# processed_data_dir2 = '/home/dihe/Projects/data/raw_data/ann_data_roberta-base-fast-trainqueryeval2_512'
# checkpoint_path2 ='/home/dihe/Projects/data/raw_data/eval_exp_21_04_14_01_1000000/ann_data_sample20q/'
# query_emb_num=4

# checkpoint_path ='/home/dihe/cudnn_file/recommender_shuqi/MIND_data/raw_data/exp_21_04_14_01/check/ann_data1000000/'
# raw_data_dir = '/home/dihe/Projects/data/raw_data/'
# processed_data_dir = '/home/dihe/Projects/data/raw_data/ann_data_roberta-base-fast-docdev_512'

# processed_data_dir2 = '/home/dihe/Projects/data/raw_data/ann_data_roberta-base-fast-trainqueryeval2_512'
# checkpoint_path2 ='/home/dihe/Projects/data/raw_data/eval_exp_21_04_14_01_1000000/ann_data_sample20q/'
# query_emb_num=4

#-------------------------------
# checkpoint_path ='/home/dihe/cudnn_file/recommender_shuqi/MIND_data/raw_data/exp_21_04_14_01/check/ann_data1000000/'
# raw_data_dir = '/home/dihe/Projects/data/raw_data/'
# processed_data_dir = '/home/dihe/Projects/data/raw_data/ann_data_roberta-base-fast-docdev_512'

# processed_data_dir2 = '/home/dihe/Projects/data/raw_data/ann_data_roberta-base-fast-trainqueryeval2_512'
# checkpoint_path2 ='/home/dihe/Projects/data/raw_data/eval_exp_21_04_14_01_1000000/ann_data_sample20q/'
# query_emb_num=4

# checkpoint_path ='/home/dihe/cudnn_file/recommender_shuqi/MIND_data/raw_data/exp_21_04_14_01/check3/ann_data30000/'
# checkpoint_path2 ='/home/dihe/Projects/data/raw_data/exp_21_04_14_01_check3_30000/ann_data_sample20q/'

# checkpoint_path ='/home/dihe/cudnn_file/recommender_shuqi/MIND_data/raw_data/exp_21_04_14_01/check3/ann_data30000/'
# raw_data_dir = '/home/dihe/Projects/data/raw_data/'
# processed_data_dir = '/home/dihe/Projects/data/raw_data/ann_data_roberta-base-fast-docdev_512'

# processed_data_dir2 = '/home/dihe/Projects/data/raw_data/ann_data_roberta-base-fast-trainqueryeval2_512'
# checkpoint_path2 ='/home/dihe/Projects/data/raw_data/exp_21_04_14_01_check3_30000/ann_data_sample20q/'
# query_emb_num=4

checkpoint_path ='/home/dihe/cudnn_file/recommender_shuqi/MIND_data/raw_data/exp_21_05_28_01/check3/ann_data10000/'
raw_data_dir = '/home/dihe/Projects/data/raw_data/'
processed_data_dir = '/home/dihe/Projects/data/raw_data/ann_data_roberta-base-fast-docdev_512'

processed_data_dir2 = '/home/dihe/Projects/data/raw_data/ann_data_roberta-base-fast-trainqueryeval2_512'
checkpoint_path2 ='/home/dihe/Projects/data/raw_data/exp_21_04_14_01_check3_30000/ann_data_sample20q/'
query_emb_num=4

#dev mrr
# checkpoint_path ='/home/dihe/Projects/data/raw_data/exp_01_05_09/ann_data910000/'
# checkpoint_path ='/home/dihe/cudnn_file/recommender_shuqi/MIND_data/raw_data/exp_21_05_21_01/check/ann_data280000/'
# raw_data_dir = '/home/dihe/Projects/data/raw_data/'
# processed_data_dir = '/home/dihe/Projects/data/raw_data/ann_data_roberta-base-fast-docdev_512'

# processed_data_dir2 = processed_data_dir
# checkpoint_path2 =checkpoint_path
# query_emb_num=8

# checkpoint_path ='/home/dihe/Projects/data/raw_data/exp_01_05_09/ann_data1000000/'
# raw_data_dir = '/home/dihe/Projects/data/raw_data/'
# processed_data_dir = '/home/dihe/Projects/data/raw_data/ann_data_roberta-base-fast-docdev_512'

# processed_data_dir2 = processed_data_dir
# checkpoint_path2 =checkpoint_path
# query_emb_num=8

# checkpoint_path ='/home/dihe/Projects/data/raw_data/exp_21_05_21_01_check/ann_data10000/'
# raw_data_dir = '/home/dihe/Projects/data/raw_data/'
# processed_data_dir = '/home/dihe/Projects/data/raw_data/ann_data_roberta-base-fast-docdev_512'

# processed_data_dir2 = processed_data_dir
# checkpoint_path2 =checkpoint_path
# query_emb_num=8

# checkpoint_path ='/home/dihe/cudnn_file/recommender_shuqi/MIND_data/raw_data/exp_21_05_21_01/check/ann_data300000/'
# raw_data_dir = '/home/dihe/Projects/data/raw_data/'
# processed_data_dir = '/home/dihe/Projects/data/raw_data/ann_data_roberta-base-fast-docdev_512'

# processed_data_dir2 = processed_data_dir
# checkpoint_path2 =checkpoint_path
# query_emb_num=8

# checkpoint_path ='/home/dihe/cudnn_file/recommender_shuqi/MIND_data/raw_data/exp_21_04_14_01/check/ann_data1000000/'
# raw_data_dir = '/home/dihe/Projects/data/raw_data/'
# processed_data_dir = '/home/dihe/Projects/data/raw_data/ann_data_roberta-base-fast-docdev_512'

# processed_data_dir2 = processed_data_dir
# checkpoint_path2 =checkpoint_path
# query_emb_num=8

# checkpoint_path ='/home/dihe/cudnn_file/recommender_shuqi/MIND_data/raw_data/exp_21_04_14_01/check3/ann_data30000/'
# raw_data_dir = '/home/dihe/Projects/data/raw_data/'
# processed_data_dir = '/home/dihe/Projects/data/raw_data/ann_data_roberta-base-fast-docdev_512'

# processed_data_dir2 = processed_data_dir
# checkpoint_path2 =checkpoint_path
# query_emb_num=8

# checkpoint_path ='/home/dihe/cudnn_file/recommender_shuqi/MIND_data/raw_data/exp_21_05_28_01/check3/ann_data10000/'
# raw_data_dir = '/home/dihe/Projects/data/raw_data/'
# processed_data_dir = '/home/dihe/Projects/data/raw_data/ann_data_roberta-base-fast-docdev_512'

# processed_data_dir2 = processed_data_dir
# checkpoint_path2 =checkpoint_path
# query_emb_num=8

#sample20
# checkpoint_path ='/home/dihe/Projects/data/raw_data/exp_01_05_09/ann_data910000/'
# raw_data_dir = '/home/dihe/Projects/data/raw_data/'
# processed_data_dir = '/home/dihe/Projects/data/raw_data/ann_data_roberta-base-fast-docdev_512'
# processed_data_dir2 = '/home/dihe/Projects/data/raw_data/ann_data_roberta-base-fast-trainqueryeval_512'
# checkpoint_path2 ='/home/dihe/Projects/data/raw_data/eval_exp_01_05_09_910000/ann_data/'
# query_emb_num=4
# processed_data_dir_origin='/home/dihe/Projects/data/raw_data/ann_data_roberta-base-fast-doc_512/'
# processed_data_dir_query_origin='/home/dihe/Projects/data/raw_data/ann_data_roberta-base-fast-doc_512/'
# ann_path='/home/dihe/Projects/data/raw_data/exp_01_05_09/ann_training_data_0'

# checkpoint_path ='/home/dihe/Projects/data/raw_data/exp_01_05_09/ann_data910000/'
# raw_data_dir = '/home/dihe/Projects/data/raw_data/'
# processed_data_dir = '/home/dihe/Projects/data/raw_data/ann_data_roberta-base-fast-docdev_512'
# processed_data_dir2 = '/home/dihe/Projects/data/raw_data/ann_data_roberta-base-fast-trainqueryeval2_512'
# checkpoint_path2 ='/home/dihe/Projects/data/raw_data/eval_exp_01_05_09_910000/ann_data_sample20q/'
# query_emb_num=4
# processed_data_dir_origin='/home/dihe/Projects/data/raw_data/ann_data_roberta-base-fast-doc_512/'
# processed_data_dir_query_origin='/home/dihe/Projects/data/raw_data/ann_data_roberta-base-fast-doc_512/'
# ann_path='/home/dihe/Projects/data/raw_data/exp_01_05_09/ann_training_data_0'

# checkpoint_path ='/home/dihe/Projects/data/raw_data/exp_21_05_21_01_check/ann_data100000/'
# raw_data_dir = '/home/dihe/Projects/data/raw_data/'
# processed_data_dir = '/home/dihe/Projects/data/raw_data/ann_data_roberta-base-fast-docdev_512'
# processed_data_dir2 = '/home/dihe/Projects/data/raw_data/ann_data_roberta-base-fast-trainqueryeval2_512'
# checkpoint_path2 ='/home/dihe/Projects/data/raw_data/exp_21_05_21_01_check_100000/ann_data_sample20q/'
# query_emb_num=4
# processed_data_dir_origin='/home/dihe/Projects/data/raw_data/ann_data_roberta-base-fast-doc_512/'
# processed_data_dir_query_origin='/home/dihe/Projects/data/raw_data/ann_data_roberta-base-fast-doc_512/'
# ann_path='/home/dihe/Projects/data/raw_data/exp_21_05_21_01_check/ann_training_data_0'

# ann_path='/home/dihe/Projects/data/raw_data/exp_21_05_21_01_check/ann_training_data_0'
# checkpoint_path ='/home/dihe/cudnn_file/recommender_shuqi/MIND_data/raw_data/exp_21_05_21_01/check/ann_data10000/'
# checkpoint_path2 ='/home/dihe/Projects/data/raw_data/exp_21_05_21_01_check_10000/ann_data_sample20q/'
# # checkpoint_path ='/home/dihe/cudnn_file/recommender_shuqi/MIND_data/raw_data/exp_01_05_09/ann_data900000/'
# # checkpoint_path2 ='/home/dihe/Projects/data/raw_data/eval_exp_01_05_09_900000/ann_data_sample20q/'

# raw_data_dir = '/home/dihe/Projects/data/raw_data/'
# processed_data_dir = '/home/dihe/Projects/data/raw_data/ann_data_roberta-base-fast-docdev_512'
# processed_data_dir2 = '/home/dihe/Projects/data/raw_data/ann_data_roberta-base-fast-trainqueryeval2_512'
# query_emb_num=4
# processed_data_dir_origin='/home/dihe/Projects/data/raw_data/ann_data_roberta-base-fast-doc_512/'
# processed_data_dir_query_origin='/home/dihe/Projects/data/raw_data/ann_data_roberta-base-fast-doc_512/'

# ann_path='/home/dihe/Projects/data/raw_data/exp_01_05_09/ann_training_data_0'
# # checkpoint_path ='/home/dihe/cudnn_file/recommender_shuqi/MIND_data/raw_data/exp_21_04_14_01/check/ann_data1000000/'
# # checkpoint_path2 ='/home/dihe/Projects/data/raw_data/eval_exp_21_04_14_01_1000000/ann_data_sample20q/'
# checkpoint_path ='/home/dihe/cudnn_file/recommender_shuqi/MIND_data/raw_data/exp_01_05_09/ann_data900000/'
# checkpoint_path2 ='/home/dihe/Projects/data/raw_data/eval_exp_01_05_09_900000/ann_data_sample20q/'
# raw_data_dir = '/home/dihe/Projects/data/raw_data/'
# processed_data_dir = '/home/dihe/Projects/data/raw_data/ann_data_roberta-base-fast-docdev_512'
# processed_data_dir2 = '/home/dihe/Projects/data/raw_data/ann_data_roberta-base-fast-trainqueryeval2_512'

# query_emb_num=4
# processed_data_dir_origin='/home/dihe/Projects/data/raw_data/ann_data_roberta-base-fast-doc_512/'
# processed_data_dir_query_origin='/home/dihe/Projects/data/raw_data/ann_data_roberta-base-fast-doc_512/'

# ann_path='/home/dihe/Projects/data/raw_data/exp_01_05_09/ann_training_data_0'
# # checkpoint_path ='/home/dihe/cudnn_file/recommender_shuqi/MIND_data/raw_data/exp_21_04_14_01/check/ann_data1000000/'
# # checkpoint_path2 ='/home/dihe/Projects/data/raw_data/eval_exp_21_04_14_01_1000000/ann_data_sample20q/'
# checkpoint_path ='/home/dihe/cudnn_file/recommender_shuqi/MIND_data/raw_data/exp_21_04_14_01/check3/ann_data30000/'
# checkpoint_path2 ='/home/dihe/Projects/data/raw_data/exp_21_04_14_01_check3_30000/ann_data_sample20q/'
# raw_data_dir = '/home/dihe/Projects/data/raw_data/'
# processed_data_dir = '/home/dihe/Projects/data/raw_data/ann_data_roberta-base-fast-docdev_512'
# processed_data_dir2 = '/home/dihe/Projects/data/raw_data/ann_data_roberta-base-fast-trainqueryeval2_512'

# query_emb_num=4
# processed_data_dir_origin='/home/dihe/Projects/data/raw_data/ann_data_roberta-base-fast-doc_512/'
# processed_data_dir_query_origin='/home/dihe/Projects/data/raw_data/ann_data_roberta-base-fast-doc_512/'

# ann_path='/home/dihe/cudnn_file/recommender_shuqi/MIND_data/raw_data/exp_21_05_28_01/ann_data/ann_training_data_0'
# checkpoint_path ='/home/dihe/cudnn_file/recommender_shuqi/MIND_data/raw_data/exp_21_05_28_01/check3/ann_data10000/'
# checkpoint_path2 ='/home/dihe/Projects/data/raw_data/exp_21_05_28_01_check3_10000/ann_data_sample20q/'
# raw_data_dir = '/home/dihe/Projects/data/raw_data/'
# processed_data_dir = '/home/dihe/Projects/data/raw_data/ann_data_roberta-base-fast-docdev_512'
# processed_data_dir2 = '/home/dihe/Projects/data/raw_data/ann_data_roberta-base-fast-trainqueryeval2_512'

# query_emb_num=4
# processed_data_dir_origin='/home/dihe/Projects/data/raw_data/ann_data_roberta-base-fast-doc_512/'
# processed_data_dir_query_origin='/home/dihe/Projects/data/raw_data/ann_data_roberta-base-fast-doc_512/'

#sample200
# checkpoint_path ='/home/dihe/Projects/data/raw_data/exp_01_05_09/ann_data910000/'
# raw_data_dir = '/home/dihe/Projects/data/raw_data/'
# processed_data_dir = '/home/dihe/Projects/data/raw_data/ann_data_roberta-base-fast-docdev_512'

# processed_data_dir2 = '/home/dihe/Projects/data/raw_data/ann_data_roberta-base-fast-trainqueryeval_512'
# checkpoint_path2 ='/home/dihe/Projects/data/raw_data/eval_exp_01_05_09_910000/ann_data/'
# query_emb_num=4

# processed_data_dir_query_origin=processed_data_dir2
# checkpoint_path_origin='/home/dihe/Projects/data/raw_data/exp_01_05_09/ann_data900000/'
# checkpoint_path_query_origin='/home/dihe/Projects/data/raw_data/eval_exp_01_05_09_900000/ann_data/'

# checkpoint_path ='/home/dihe/Projects/data/raw_data/exp_01_05_09/ann_data910000/'
# raw_data_dir = '/home/dihe/Projects/data/raw_data/'
# processed_data_dir = '/home/dihe/Projects/data/raw_data/ann_data_roberta-base-fast-docdev_512'

# processed_data_dir2 = '/home/dihe/Projects/data/raw_data/ann_data_roberta-base-fast-trainqueryeval2_512'
# checkpoint_path2 ='/home/dihe/Projects/data/raw_data/eval_exp_01_05_09_910000/ann_data_sample20q/'
# query_emb_num=4

# processed_data_dir_query_origin=processed_data_dir2
# checkpoint_path_origin='/home/dihe/Projects/data/raw_data/exp_01_05_09/ann_data900000/'
# checkpoint_path_query_origin='/home/dihe/Projects/data/raw_data/eval_exp_01_05_09_900000/ann_data_sample20q/'

# checkpoint_path ='/home/dihe/Projects/data/raw_data/exp_21_05_21_01_check/ann_data100000/'
# raw_data_dir = '/home/dihe/Projects/data/raw_data/'
# processed_data_dir = '/home/dihe/Projects/data/raw_data/ann_data_roberta-base-fast-docdev_512'

# processed_data_dir2 = '/home/dihe/Projects/data/raw_data/ann_data_roberta-base-fast-trainqueryeval2_512'
# checkpoint_path2 ='/home/dihe/Projects/data/raw_data/exp_21_05_21_01_check_100000/ann_data_sample20q/'
# query_emb_num=4

# processed_data_dir_query_origin=processed_data_dir2
# checkpoint_path_origin='/home/dihe/Projects/data/raw_data/exp_01_05_09/ann_data900000/'
# checkpoint_path_query_origin='/home/dihe/Projects/data/raw_data/eval_exp_01_05_09_900000/ann_data_sample20q/'

# checkpoint_path ='/home/dihe/cudnn_file/recommender_shuqi/MIND_data/raw_data/exp_21_04_14_01/check/ann_data1000000/'
# raw_data_dir = '/home/dihe/Projects/data/raw_data/'
# processed_data_dir = '/home/dihe/Projects/data/raw_data/ann_data_roberta-base-fast-docdev_512'

# processed_data_dir2 = '/home/dihe/Projects/data/raw_data/ann_data_roberta-base-fast-trainqueryeval2_512'
# checkpoint_path2 ='/home/dihe/Projects/data/raw_data/eval_exp_21_04_14_01_1000000/ann_data_sample20q/'
# query_emb_num=4

# processed_data_dir_query_origin=processed_data_dir2
# checkpoint_path_origin='/home/dihe/Projects/data/raw_data/exp_01_05_09/ann_data900000/'
# checkpoint_path_query_origin='/home/dihe/Projects/data/raw_data/eval_exp_01_05_09_900000/ann_data_sample20q/'

# checkpoint_path ='/home/dihe/cudnn_file/recommender_shuqi/MIND_data/raw_data/exp_21_04_14_01/check3/ann_data30000/'
# raw_data_dir = '/home/dihe/Projects/data/raw_data/'
# processed_data_dir = '/home/dihe/Projects/data/raw_data/ann_data_roberta-base-fast-docdev_512'

# processed_data_dir2 = '/home/dihe/Projects/data/raw_data/ann_data_roberta-base-fast-trainqueryeval2_512'
# checkpoint_path2 ='/home/dihe/Projects/data/raw_data/exp_21_04_14_01_check3_30000/ann_data_sample20q/'
# query_emb_num=4

# processed_data_dir_query_origin=processed_data_dir2
# checkpoint_path_origin='/home/dihe/Projects/data/raw_data/exp_01_05_09/ann_data900000/'
# checkpoint_path_query_origin='/home/dihe/Projects/data/raw_data/eval_exp_01_05_09_900000/ann_data_sample20q/'

# checkpoint_path ='/home/dihe/cudnn_file/recommender_shuqi/MIND_data/raw_data/exp_21_05_28_01/check3/ann_data10000/'
# raw_data_dir = '/home/dihe/Projects/data/raw_data/'
# processed_data_dir = '/home/dihe/Projects/data/raw_data/ann_data_roberta-base-fast-docdev_512'

# processed_data_dir2 = '/home/dihe/Projects/data/raw_data/ann_data_roberta-base-fast-trainqueryeval2_512'
# checkpoint_path2 ='/home/dihe/Projects/data/raw_data/exp_21_05_28_01_check3_10000/ann_data_sample20q/'
# query_emb_num=4

# processed_data_dir_query_origin=processed_data_dir2
# checkpoint_path_origin='/home/dihe/cudnn_file/recommender_shuqi/MIND_data/raw_data/exp_21_04_14_01/check3/ann_data30000/'
# checkpoint_path_query_origin='/home/dihe/Projects/data/raw_data/exp_21_04_14_01_check3_30000/ann_data_sample20q/'

if data_type == 0:
    topN = 100
else:
    topN = 1000

qidmap_path = processed_data_dir2+"/qid2offset.pickle"
pidmap_path = processed_data_dir+"/pid2offset.pickle"
# if data_type == 0:
#     if test_set == 1:
#         query_path = raw_data_dir+"/docleaderboard-queries.tsv"
#         passage_path = raw_data_dir+"/docleaderboard-top100.tsv"
#     else:
#         query_path = raw_data_dir+"/msmarco-docdev-queries.tsv"
#         passage_path = raw_data_dir+"/msmarco-docdev-top100"
# else:
#     if test_set == 1:
#         query_path = raw_data_dir+"/msmarco-test2019-queries.tsv"
#         passage_path = raw_data_dir+"/msmarco-passagetest2019-top1000.tsv"
#     else:
#         query_path = raw_data_dir+"/queries.dev.small.tsv"
#         passage_path = raw_data_dir+"/top1000.dev.tsv"

with open(qidmap_path, 'rb') as handle:
    qidmap = pickle.load(handle)

with open(pidmap_path, 'rb') as handle:
    pidmap = pickle.load(handle)

pidmap_re={}
for item in pidmap:
    assert pidmap[item] not in pidmap_re
    pidmap_re[pidmap[item]]=item #'D'+str(item)

qidmap_re={}
for item in qidmap:
    assert qidmap[item] not in qidmap_re
    qidmap_re[qidmap[item]]=item

def get_reverse_dict(mydict):
    mydict_re={}
    print(mydict)
    for item in mydict:
        assert mydict[item] not in mydict_re
        mydict_re[mydict[item]]=item
    return mydict_re

count_none=0
dev_query_positive_id = {}
query_positive_id_path = os.path.join(raw_data_dir, "msmarco-doctrain-qrels.tsv")
#query_positive_id_path = os.path.join(raw_data_dir, "msmarco-docdev-qrels.tsv")

with open(query_positive_id_path, 'r', encoding='utf8') as f:
    tsvreader = csv.reader(f, delimiter=" ")
    for [topicid,_, docid, rel] in tsvreader:
        topicid = int(topicid)
        docid = int(docid[1:])
        if topicid not in dev_query_positive_id:
            dev_query_positive_id[topicid] = {}
        dev_query_positive_id[topicid][docid] = int(rel)
        assert len(dev_query_positive_id[topicid])==1
        if docid not in pidmap:
            count_none+=1
print('count_none: ',count_none)

# qset = set()
# with gzip.open(query_path, 'rt', encoding='utf-8') if query_path[-2:] == "gz" else open(query_path, 'rt', encoding='utf-8') as f:
#     tsvreader = csv.reader(f, delimiter="\t")
#     for [qid, query] in tsvreader:
#         qset.add(qid)

# bm25 = collections.defaultdict(set)
# with gzip.open(passage_path, 'rt', encoding='utf-8') if passage_path[-2:] == "gz" else open(passage_path, 'rt', encoding='utf-8') as f:
#     for line in tqdm(f):
#         if data_type == 0:
#             [qid, Q0, pid, rank, score, runstring] = line.split(' ')
#             pid = pid[1:]
#         else:
#             [qid, pid, query, passage] = line.split("\t")
#             #print('???',qid)
#         if qid in qset and int(qid) in qidmap:
#             bm25[qidmap[int(qid)]].add(pidmap[int(pid)])
#         # else:
#         #     print('???',qid,qid in qset)

# #assert 1==0
# print("number of queries with " +str(topN) + " BM25 passages:", len(bm25))

def get_sample20():

    #train_queries={}
    with open(processed_data_dir_query_origin+"/qid2offset_train.pickle", 'rb') as handle:
        qidmap_origin = pickle.load(handle)

    with open(processed_data_dir_query_origin+"/pid2offset.pickle", 'rb') as handle:
        pidmap_origin = pickle.load(handle)
    qidmap_origin_re={}
    for item in qidmap_origin:
        assert qidmap_origin[item] not in qidmap_origin_re
        qidmap_origin_re[qidmap_origin[item]]=item

    pidmap_origin_re={}
    for item in pidmap_origin:
        assert pidmap_origin[item] not in pidmap_origin_re
        pidmap_origin_re[pidmap_origin[item]]=item

    train_q_sample20={}
    with open(ann_path, 'r') as f:
        ann_training_data = f.readlines()

    # aligned_size = (len(ann_training_data) // 8) * 8
    # ann_training_data = ann_training_data[:aligned_size]

    # passage_embedding2id=[]
    # for i in range(8):
    #     # with open(checkpoint_path + "passage_"+str(checkpoint)+"__emb_p__data_obj_"+str(i)+".pb", 'rb') as handle:
    #     #     passage_embedding.append(pickle.load(handle))
    #     #     print('ok???',3,i)
    #     with open(checkpoint_path_origin + "passage_"+str(checkpoint)+"__embid_p__data_obj_"+str(i)+".pb", 'rb') as handle:
    #         passage_embedding2id.append(pickle.load(handle))
    #         print('ok???',4,i)

    # #passage_embedding2id = np.concatenate(passage_embedding2id, axis=0)

    # dev_query_embedding2id=[]
    # for i in range(4):
    #     print('???',checkpoint_path_query_origin + "dev_query_"+str(checkpoint)+"__emb_p__data_obj_"+str(i)+".pb")
    #     #with open(checkpoint_path2 + "dev_query_"+str(checkpoint)+"__emb_p__data_obj_"+str(i)+".pb", 'rb') as handle:
    #     with open(checkpoint_path_query_origin + "dev_query_"+str(checkpoint)+"__embid_p__data_obj_"+str(i)+".pb", 'rb') as handle:
    #         dev_query_embedding2id.append(pickle.load(handle))
    #         print('ok???',2)

    # if (not dev_query_embedding2id) or not (passage_embedding2id):
    #     print("No data found for checkpoint: ",checkpoint)

    # passage_embedding2id = np.concatenate(passage_embedding2id, axis=0)
    # dev_query_embedding2id=np.concatenate(dev_query_embedding2id, axis=0)

    # passage_embedding2id2_r=reverse_dict(passage_embedding2id2)
    # dev_query_embedding2id_r=reverse_dict(dev_query_embedding2id)

    for line in ann_training_data:
        line_arr=line.strip().split('\t')
        # qid = qidmap_origin_re[dev_query_embedding2id[int(line_arr[0])]]
        # if qid in dev_query_positive_id:
        #     pos_pid = pidmap_re[passage_embedding2id[int(line_arr[1])]]
        #     neg_pids = line_arr[2].split(',')
        #     neg_pids = [pidmap_re[passage_embedding2id[int(neg_pid)]] for neg_pid in neg_pids]
        #     train_q_sample20[qid]=neg_pids+[pos_pid]
        qid = qidmap_origin_re[int(line_arr[0])]
        if qid in dev_query_positive_id:
            pos_pid = pidmap_origin_re[int(line_arr[1])]
            neg_pids = line_arr[2].split(',')
            neg_pids = [pidmap_origin_re[int(neg_pid)] for neg_pid in neg_pids]
            train_q_sample20[qid]=neg_pids[:20]+[pos_pid]
        else:
            assert 1==0

    return train_q_sample20

def get_sample200():

    #dev_query_positive_id[topicid]

    with open(processed_data_dir_query_origin+"/qid2offset.pickle", 'rb') as handle:
        qidmap_origin = pickle.load(handle)
    qidmap_origin_re={}
    for item in qidmap_origin:
        assert qidmap_origin[item] not in qidmap_origin_re
        qidmap_origin_re[qidmap_origin[item]]=item

    dev_query_embedding=[]
    dev_query_embedding2id=[]
    train_q_sample200={}
    for i in range(4):
        #try:
        print('???',checkpoint_path_query_origin + "dev_query_"+str(checkpoint)+"__emb_p__data_obj_"+str(i)+".pb")
        with open(checkpoint_path_query_origin + "dev_query_"+str(checkpoint)+"__emb_p__data_obj_"+str(i)+".pb", 'rb') as handle:
            dev_query_embedding.append(pickle.load(handle))
            print('ok1???')
        with open(checkpoint_path_query_origin + "dev_query_"+str(checkpoint)+"__embid_p__data_obj_"+str(i)+".pb", 'rb') as handle:
            dev_query_embedding2id.append(pickle.load(handle))
            print('ok???',2)

    passage_embedding=[]
    passage_embedding2id=[]
    for i in range(8):
        with open(checkpoint_path_origin + "passage_"+str(checkpoint)+"__emb_p__data_obj_"+str(i)+".pb", 'rb') as handle:
            passage_embedding.append(pickle.load(handle))
            print('ok???',3,i)
        with open(checkpoint_path_origin + "passage_"+str(checkpoint)+"__embid_p__data_obj_"+str(i)+".pb", 'rb') as handle:
            passage_embedding2id.append(pickle.load(handle))
            print('ok???',4,i)

    if (not dev_query_embedding) or (not dev_query_embedding2id) or (not passage_embedding) or not (passage_embedding2id):
        print("No data found for checkpoint: ",checkpoint)

    dev_query_embedding = np.concatenate(dev_query_embedding, axis=0)
    dev_query_embedding2id = np.concatenate(dev_query_embedding2id, axis=0)
    passage_embedding = np.concatenate(passage_embedding, axis=0)
    passage_embedding2id = np.concatenate(passage_embedding2id, axis=0)

    dim = passage_embedding.shape[1]
    faiss.omp_set_num_threads(16)
    cpu_index = faiss.IndexFlatIP(dim)
    cpu_index.add(passage_embedding)
    _, dev_I = cpu_index.search(dev_query_embedding, 200)

    #dev_query_embedding2id_r=get_reverse_dict(dev_query_embedding2id)

    #for item in train_queries:
    query_list=[]
    # for i,query_idx in enumerate(dev_query_embedding2id):
    for i,query_idx in enumerate(range(len(dev_I))):
        query_id = qidmap_origin_re[dev_query_embedding2id[query_idx]]
        selected_ann_idx=dev_I[query_idx]
        if query_id in dev_query_positive_id:
            train_q_sample200[query_id]=[]
            pos_id=list(dev_query_positive_id[query_id].keys())[0]
            for idx in selected_ann_idx:
                pred_pid = pidmap_re[passage_embedding2id[idx]]
                train_q_sample200[query_id].append(pred_pid)

            train_q_sample200[query_id]=train_q_sample200[query_id][:20]

            if pos_id not in  train_q_sample200[query_id]:
                train_q_sample200[query_id]+=[pos_id]
        # if i<5:
        #     print(query_id,dev_query_embedding2id[query_idx],query_idx)
        #     query_list.append(query_id)

    #print([train_q_sample200[x] for x in query_list])

    return train_q_sample200

def get_all(passage_embedding,passage_embedding2id):

    for i in range(8):
        #try:
        # print('???',checkpoint_path2 + "dev_query_"+str(checkpoint)+"__emb_p__data_obj_"+str(i)+".pb")
        # with open(checkpoint_path2 + "dev_query_"+str(checkpoint)+"__emb_p__data_obj_"+str(i)+".pb", 'rb') as handle:
        #     dev_query_embedding.append(pickle.load(handle))
        #     print('ok1???')
        # with open(checkpoint_path2 + "dev_query_"+str(checkpoint)+"__embid_p__data_obj_"+str(i)+".pb", 'rb') as handle:
        #     dev_query_embedding2id.append(pickle.load(handle))
        #     print('ok???',2)
        with open(checkpoint_path + "passage_"+str(checkpoint)+"__emb_p__data_obj_"+str(i)+".pb", 'rb') as handle:
            passage_embedding.append(pickle.load(handle))
            print('ok???',3,i)
        with open(checkpoint_path + "passage_"+str(checkpoint)+"__embid_p__data_obj_"+str(i)+".pb", 'rb') as handle:
            passage_embedding2id.append(pickle.load(handle))
            print('ok???',4,i)
        # except:
        #     break
    if (not passage_embedding) or not (passage_embedding2id):
        print("No data found for checkpoint: ",checkpoint)
    passage_embedding = np.concatenate(passage_embedding, axis=0)
    passage_embedding2id = np.concatenate(passage_embedding2id, axis=0)
    return passage_embedding,passage_embedding2id

def convert_to_string_id(result_dict):
    string_id_dict = {}

    # format [string, dict[string, val]]
    for k, v in result_dict.items():
        _temp_v = {}
        for inner_k, inner_v in v.items():
            _temp_v[str(inner_k)] = inner_v

        string_id_dict[str(k)] = _temp_v

    return string_id_dict

def EvalDevQuery(query_embedding2id, passage_embedding2id, qidmap_re,pidmap_re, dev_query_positive_id,I_nearest_neighbor,topN,bm25=None):
    prediction = {} #[qid][docid] = docscore, here we use -rank as score, so the higher the rank (1 > 2), the higher the score (-1 > -2)

    #w=open('result_eval.txt','w')
    total = 0
    labeled = 0
    Atotal = 0
    Alabeled = 0
    qids_to_ranked_candidate_passages = {}
    mrr=0.0
    mycount=0
    for query_idx in range(len(I_nearest_neighbor)):
        seen_pid = set()
        query_id = qidmap_re[query_embedding2id[query_idx]]

        if bm25 and query_id not in bm25:
            #assert 1==0
            continue

        prediction[query_id] = {}

        top_ann_pid = I_nearest_neighbor[query_idx].copy()
        selected_ann_idx = top_ann_pid[:topN]
        #print('???',topN)

        #if train_q_sample20 !=None:

        rank = 0
        flag=0

        if query_id in qids_to_ranked_candidate_passages:
            assert 1==0,"query not in"
            pass
        else:
            # By default, all PIDs in the list of 1000 are 0. Only override those that are given
            tmp = [0] * 1000
            qids_to_ranked_candidate_passages[query_id] = tmp

        mycount+=1
        for idx in selected_ann_idx:
            pred_pid = pidmap_re[passage_embedding2id[idx]]

            if not pred_pid in seen_pid:
                # this check handles multiple vector per document
                qids_to_ranked_candidate_passages[query_id][rank]=pred_pid
                #w.write(str(query_id)+'\t'+str(pred_pid)+'\t'+str(rank+1)+'\n')
                # assert len(dev_query_positive_id[query_id]) ==1
                # for item in dev_query_positive_id[query_id]:
                #     assert item in pidmap

                assert pred_pid in pidmap
                if pred_pid in dev_query_positive_id[query_id]:
                    mrr += 1/(rank + 1)
                    flag=1
                    #print('rank: ',rank)

                Atotal += 1
                if pred_pid not in dev_query_positive_id[query_id]:
                    Alabeled += 1
                if rank < 10:
                    total += 1
                    if pred_pid not in dev_query_positive_id[query_id]:
                        labeled += 1
                rank += 1
                prediction[query_id][pred_pid] = -rank
                seen_pid.add(pred_pid)

        #assert rank!=0, "pos not in"
    # w.close()
    evaluator = pytrec_eval.RelevanceEvaluator(
        convert_to_string_id(dev_query_positive_id), {'map_cut', 'ndcg_cut', 'recip_rank','recall'})
    eval_query_cnt = 0
    result = evaluator.evaluate(convert_to_string_id(prediction))

    print('???',mrr/mycount,mycount,mrr,len(I_nearest_neighbor))

    qids_to_relevant_passageids = {}
    for qid in dev_query_positive_id:
        qid = int(qid)
        if qid in qids_to_relevant_passageids:
            pass
        else:
            qids_to_relevant_passageids[qid] = []
            for pid in dev_query_positive_id[qid]:
                #assert pid>0
                if pid>0:
                    qids_to_relevant_passageids[qid].append(pid)

    if data_type == 0:
        MaxMRRRank=100
    else:
        MaxMRRRank=10

    ms_mrr = compute_metrics(qids_to_relevant_passageids, qids_to_ranked_candidate_passages,MaxMRRRank=MaxMRRRank)
    # ms_mrr = compute_metrics(dev_query_positive_id, qids_to_ranked_candidate_passages,MaxMRRRank=MaxMRRRank)
    # print('???',ms_mrr)

    ndcg = 0
    Map = 0
    mrr = 0
    recall = 0
    recall_1000 = 0

    for k in result.keys():
        eval_query_cnt += 1
        ndcg += result[k]["ndcg_cut_10"]
        Map += result[k]["map_cut_10"]
        mrr += result[k]["recip_rank"]
        recall += result[k]["recall_"+str(topN)]

    final_ndcg = ndcg / eval_query_cnt
    final_Map = Map / eval_query_cnt
    final_mrr = mrr / eval_query_cnt
    final_recall = recall / eval_query_cnt
    hole_rate = labeled/total
    Ahole_rate = Alabeled/Atotal

    return final_ndcg, eval_query_cnt, final_Map, final_mrr, final_recall, hole_rate, ms_mrr, Ahole_rate, result, prediction

dev_query_embedding = []
dev_query_embedding2id = []
passage_embedding = []
passage_embedding2id = []

for i in range(query_emb_num):
    #try:
    print('???',checkpoint_path2 + "dev_query_"+str(checkpoint)+"__emb_p__data_obj_"+str(i)+".pb")
    with open(checkpoint_path2 + "dev_query_"+str(checkpoint)+"__emb_p__data_obj_"+str(i)+".pb", 'rb') as handle:
        dev_query_embedding.append(pickle.load(handle))
        print('ok1???')
    with open(checkpoint_path2 + "dev_query_"+str(checkpoint)+"__embid_p__data_obj_"+str(i)+".pb", 'rb') as handle:
        dev_query_embedding2id.append(pickle.load(handle))
        print('ok???',2)

passage_embedding,passage_embedding2id=get_all(passage_embedding,passage_embedding2id)

# sample20 =get_sample200(passage_embedding)

# passage_embedding,passage_embedding2id =get_sample200(passage_embedding)

if (not dev_query_embedding) or (not dev_query_embedding2id):
    print("No data found for checkpoint: ",checkpoint)

dev_query_embedding = np.concatenate(dev_query_embedding, axis=0)
dev_query_embedding2id = np.concatenate(dev_query_embedding2id, axis=0)

##reranking
#sample20 =get_sample20()

# sample20 =get_sample200()
# pidmap_t = collections.defaultdict(list)
# for i in range(len(passage_embedding2id)):
#     pidmap_t[pidmap_re[passage_embedding2id[i]]].append(i)  # abs pos(key) to rele pos(val)
# all_dev_I = []
# for i,qid in enumerate(range(len(dev_query_embedding2id))):
#     qid_r=qidmap_re[dev_query_embedding2id[qid]]
#     p_set = []
#     p_set_map = {}
#     if qid_r not in sample20:
#         print('no')

#     else:
#         #print('yes')
#         count = 0
#         for k,pid in enumerate(sample20[qid_r]):
#             if pid in pidmap_t:
#                 for val in pidmap_t[pid]:
#                     p_set.append(passage_embedding[val])
#                     p_set_map[count] = val # new rele pos(key) to old rele pos(val)
#                     count += 1
#             else:
#                 print(pid,"not in passages")
#     #print('???len(p_set)',len(p_set))
#     if len(p_set)==0:
#         all_dev_I.append([-1]*10)
#     else:

#         dim = passage_embedding.shape[1]
#         faiss.omp_set_num_threads(16)
#         cpu_index = faiss.IndexFlatIP(dim)
#         p_set =  np.asarray(p_set)
#         cpu_index.add(p_set)
#         _, dev_I = cpu_index.search(dev_query_embedding[i:i+1], len(p_set))
#         # if i<5:
#         #     print(sample20[qid_r],qid_r,dev_query_embedding2id[qid],qid)
#         # if i<5:
#         #     print(dev_I,dev_query_positive_id[qid_r])
#         for j in range(len(dev_I[0])):
#             dev_I[0][j] = p_set_map[dev_I[0][j]]
#         # if i<5:
#         #     print(dev_I,dev_query_positive_id[qid_r])
#         #     print([pidmap_re[passage_embedding2id[x]] for x in dev_I[0]])
#         #     print('-----------------------')
#         all_dev_I.append(dev_I[0])
# print(len(sample20),len(all_dev_I))
# result = EvalDevQuery(dev_query_embedding2id, passage_embedding2id, qidmap_re,pidmap_re, dev_query_positive_id, all_dev_I, 10,bm25=sample20)
# final_ndcg, eval_query_cnt, final_Map, final_mrr, final_recall, hole_rate, ms_mrr, Ahole_rate, metrics, prediction = result
# print("Reranking Results for checkpoint "+str(checkpoint))
# print("Reranking NDCG@10:" + str(final_ndcg))
# print("Reranking map@10:" + str(final_Map))
# print("Reranking pytrec_mrr:" + str(final_mrr))
# print("Reranking recall@"+str(topN)+":" + str(final_recall))
# print("Reranking hole rate@10:" + str(hole_rate))
# print("Reranking hole rate:" + str(Ahole_rate))
# print("Reranking ms_mrr:" + str(ms_mrr))

#full ranking
dim = passage_embedding.shape[1]
faiss.omp_set_num_threads(16)
cpu_index = faiss.IndexFlatIP(dim)
cpu_index.add(passage_embedding)
_, dev_I = cpu_index.search(dev_query_embedding, topN)
#print('???',dev_I[:10])
result = EvalDevQuery(dev_query_embedding2id, passage_embedding2id, qidmap_re,pidmap_re , dev_query_positive_id,dev_I, 10)

final_ndcg, eval_query_cnt, final_Map, final_mrr, final_recall, hole_rate, ms_mrr, Ahole_rate, metrics, prediction = result
print("Results for checkpoint "+str(checkpoint))
print("NDCG@10:" + str(final_ndcg))
print("map@10:" + str(final_Map))
print("pytrec_mrr:" + str(final_mrr))
print("recall@"+str(topN)+":" + str(final_recall))
print("hole rate@10:" + str(hole_rate))
print("hole rate:" + str(Ahole_rate))
print("ms_mrr:" + str(ms_mrr))


text:
import numpy as np
# Some physical constants at sea level
R0 = 6371.0008 # volumetric radius of the Earth, [km]
g0 = 9.80665 # standard gravity acceleration, [m/s^2]
M0 = 28.9644 # mean molecular weight, [kg/kmol]
R_air = 8314.32/M0 # gas constant for dry air, [J/kg/K]
gamma = 1.4 # specific heat ratio, namely the ratio of constant-pressure to constant-volume specific heat
T0 = 288.15 # temperature, [K]
p0 = 101325 # pressure, [Pa]
h0 = 0 # height, [km]

pi = np.pi
twopi = 2*np.pi
fourpi = 4*np.pi
degrad = pi / 180
al10 = np.log(10)

# AVOGAD is Avogadro's number in mks units (molecules/kmol)
avogad = 6.02257e26
pivo2 = 1.5707963
pivo4 = 1.5707963

# RSTAR is the universal gas-constant in mks units (joules/K/kmol)
rstar = 8314.32

x = np.arange(0.5,48)/24


text:
import argparse
import random
import os
import json

import numpy as np
import torch
import torch.optim as optim
import visdom
from tqdm import tqdm

# from dataset import *
from data_loader import DataLoader
from model import *
from utils import *
import emd_module as emd
from chamfer import dist_chamfer_3D

chamfer_distance = dist_chamfer_3D.chamfer_3DDist()

def dist_cd(pc1, pc2):
    dist1, dist2, _, _ = chamfer_distance(pc1, pc2)
    return (torch.mean(dist1) + torch.mean(dist2)) / 2 * 100

def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument('--batch_size', type=int, default=36,
                        help='input batch size')
    parser.add_argument('--workers', type=int, default=4,
                        help='number of data loading workers')
    parser.add_argument('--data_dir', type=str, default='data/shapenet_color',
                        help='Data directory')
    parser.add_argument('--nepoch', type=int, default=25,
                        help='number of epochs to train for')
    parser.add_argument('--model', type=str, default='',
                        help='optional reload model path')
    parser.add_argument('--in_num_points', type=int, default=2048,
                        help='number of output points')
    parser.add_argument('--out_num_points', type=int, default=4096,
                        help='number of input points')
    parser.add_argument('--n_primitives', type=int, default=16,
                        help='number of surface elements')
    parser.add_argument('--exp_name', type=str, default="shapenet_02691156",
                        help='Experiment name')
    parser.add_argument('--category', type=str, default='02691156', metavar='N',
                        help='category of the data')
    parser.add_argument('--isrotate', action='store_true', default=True,
                        help='Whether to rotate the input data along Z axis [default: True]')
    parser.add_argument('--eval', action='store_true', default=False,
                        help='Evaluate the model')
    parser.add_argument('--debug', action='store_true', default=False,
                        help='Small dataset to debug')
    parser.add_argument('--multi_gpu', action='store_true', default=False,
                        help='Use multiple GPU or not')

    opt = parser.parse_args()
    print(opt)

    return opt

def train_one_epoch(network, optimizer, dataloader, train_loss, vis):
    # Log items
    network.train()
    emd_total = 0
    chamfer_total = 0
    for i, data in tqdm(enumerate(dataloader), total=len(dataloader),
                        smoothing=0.9):
        # Load data
        input1, input2, gt1, gt2 = data
        input1 = input1.float().cuda()
        input2 = input2.float().cuda()
        gt1 = gt1.float().cuda()
        gt2 = gt2.float().cuda()
        input1 = input1.transpose(2, 1).contiguous()
        input2 = input2.transpose(2, 1).contiguous()

        # Optimize for input1
        optimizer.zero_grad()
        output11, output12, emd11, emd12, expansion_penalty1 = \
            network(input1, gt1.contiguous(), 0.01, 100)
        output21, output22, emd21, emd22, expansion_penalty2 = \
            network(input2, gt2.contiguous(), 0.01, 100)
        loss_net1 = emd11.mean() + emd12.mean() + expansion_penalty1.mean() * 0.1
        loss_net2 = emd21.mean() + emd22.mean() + expansion_penalty2.mean() * 0.1
        loss_total = (loss_net1 + loss_net2) / 2
        loss_total.backward()
        optimizer.step()

        # Log
        emd_total += (emd12.mean().item() + emd22.mean().item()) / 2
        chamfer_total += (dist_cd(output12, gt1).item() +
                          dist_cd(output22, gt2).item()) / 2
        train_loss.update((emd12.mean().item() + emd22.mean().item()) / 2)

        if i % 10 == 0:
            idx = random.randint(0, input1.size()[0] - 1)
            vis.scatter(X = gt1.contiguous()[idx].data.cpu()[:, :3],
                        win = 'TRAIN_GT',
                        opts = dict(
                            title='Train GT',
                            markersize = 2,
                            ),
                        )
            vis.scatter(X = input1.transpose(2,1).contiguous()[idx].data.cpu(),
                        win = 'TRAIN_INPUT',
                        opts = dict(
                            title='Train input',
                            markersize = 2,
                            ),
                        )
            vis.scatter(X = output11[idx].data.cpu(),
                        Y = labels_generated_points[0:output11.size(1)],
                        win = 'TRAIN_COARSE',
                        opts = dict(
                            title='Train coarse',
                            markersize=2,
                            ),
                        )
            vis.scatter(X = output12[idx].data.cpu(),
                        win = 'TRAIN_OUTPUT',
                        opts = dict(
                            title='Train output',
                            markersize=2,
                            ),
                        )

    emd = emd_total / len(dataloader)
    chamfer = chamfer_total / len(dataloader)
    return emd, chamfer

def evaluate(network, dataloader, val_loss, vis, precise=False):
    network.eval()
    emd_total = 0
    chamfer_total = 0

    for i, data in tqdm(enumerate(dataloader), total=len(dataloader), smoothing=0.9):
        input1, input2, gt1, gt2 = data
        input1 = input1.float().cuda()
        input2 = input2.float().cuda()
        gt1 = gt1.float().cuda()
        gt2 = gt2.float().cuda()
        input1 = input1.transpose(2,1).contiguous()
        input2 = input2.transpose(2,1).contiguous()

        with torch.no_grad():
            # if precise:
            #     eps = 0.002
            #     n_iters = 10000
            # else:
            #     eps = 0.004
            #     n_iters = 3000
            eps = 0.01
            n_iters = 100
            output11, output12, emd11, emd12, expansion_penalty1 = \
                network(input1, gt1.contiguous(), eps, n_iters)
            output21, output22, emd21, emd22, expansion_penalty2 = \
                network(input2, gt2.contiguous(), eps, n_iters)
        val_loss.update((emd12.mean().item() + emd22.mean().item()) / 2)
        idx = random.randint(0, input1.size()[0] - 1)
        if i % 10 == 0:
            vis.scatter(X = gt1.contiguous()[idx].data.cpu()[:, :3],
                        win = 'VAL_GT',
                        opts = dict(
                            title='Val GT',
                            markersize = 2,
                        ),
            )
            vis.scatter(X = input1.transpose(2,1).contiguous()[idx].data.cpu(),
                        win = 'VAL_INPUT',
                        opts = dict(
                            title='Val input',
                            markersize = 2,
                        ),
            )
            vis.scatter(X = output11[idx].data.cpu(),
                        Y = labels_generated_points[0:output11.size(1)],
                        win = 'VAL_COARSE',
                        opts = dict(
                            title='Val coarse',
                            markersize=2,
                        ),
            )
            vis.scatter(X = output12[idx].data.cpu(),
                        win = 'VAL_OUTPUT',
                        opts = dict(
                            title='Val output',
                            markersize=2,
                        ),
            )
        emd_total += (emd12.mean().item() + emd22.mean().item()) / 2
        with torch.no_grad():
            chamfer_total += (dist_cd(output12, gt1).item() +
                              dist_cd(output22, gt2).item()) / 2

    emd = emd_total / len(dataloader)
    chamfer = chamfer_total / len(dataloader)
    return emd, chamfer

if __name__ == '__main__':
    opt = parse_args()

    class FullModel(nn.Module):
        def __init__(self, model):
            super(FullModel, self).__init__()
            self.model = model
            self.EMD = emd.emdModule()

        def forward(self, inputs, gt, eps, iters):
            output1, output2, expansion_penalty = self.model(inputs)
            gt = gt[:, :, :3]

            dist, _ = self.EMD(output1, gt, eps, iters)
            emd1 = torch.sqrt(dist).mean(1)

            dist, _ = self.EMD(output2, gt, eps, iters)
            emd2 = torch.sqrt(dist).mean(1)

            return output1, output2, emd1, emd2, expansion_penalty

    ''' Logs setup '''
    save_path = opt.exp_name
    if not os.path.exists('./logs/'):
        os.mkdir('./logs/')
    log_dir =  os.path.join('logs', save_path)
    if not os.path.exists(log_dir):
        os.mkdir(log_dir)
    logname = os.path.join(log_dir, 'log.txt')
    if not opt.eval:
        os.system('cp ./main.py %s' % log_dir)
        os.system('cp ./data_loader.py %s' % log_dir)
        os.system('cp ./model.py %s' % log_dir)

    ''' Seeding '''
    if not opt.eval:
        opt.manualSeed = random.randint(1, 10000)
        print("Random Seed: ", opt.manualSeed)
        random.seed(opt.manualSeed)
        torch.manual_seed(opt.manualSeed)

    ''' Initialize network '''
    network = MSN(num_points = opt.out_num_points,
                  n_primitives = opt.n_primitives)
    if opt.multi_gpu and torch.cuda.device_count() > 1:
        network = torch.nn.DataParallel(FullModel(network))
        network.cuda()
        model = network.module.model
        opt.batch_size *= torch.cuda.device_count()
    else:
        network = FullModel(network)
        network.cuda()
        model = network.model
    model.apply(weights_init) #initialization of the weight

    if opt.model != '':
        model.load_state_dict(torch.load(opt.model))
        print("Trained weights loaded.")
    elif opt.eval:
        model_path = os.path.join(log_dir, 'network.pth')
        if os.path.exists(model_path):
            model.load_state_dict(torch.load(model_path))
            print("Trained weights loaded.")
        else:
            print("Trained model path should be specified!")
            print("Exiting...")
            exit(0)

    lrate = 0.001 #learning rate
    optimizer = optim.Adam(model.parameters(), lr = lrate)
    print("Model initialized.")

    ''' Load data '''
    dataset = DataLoader(root=opt.data_dir,
                         npoint=opt.in_num_points,
                         split='train',
                         category=[opt.category],
                         isrotate=opt.isrotate,
                         debug=opt.debug)
    dataloader = torch.utils.data.DataLoader(dataset,
                                             batch_size=opt.batch_size,
                                             shuffle=True,
                                             num_workers=int(opt.workers),
                                             pin_memory=True,
                                             drop_last=True)
    dataset_test = DataLoader(root=opt.data_dir,
                              npoint=opt.in_num_points,
                              split='test',
                              category=[opt.category],
                              isrotate=opt.isrotate,
                              debug=opt.debug)
    dataloader_test = torch.utils.data.DataLoader(dataset_test,
                                                  batch_size=opt.batch_size,
                                                  shuffle=False,
                                                  num_workers=int(opt.workers),
                                                  pin_memory=True,
                                                  drop_last=True)

    ''' Visualization setup'''
    vis = visdom.Visdom(port=8097, env=opt.exp_name) # set your port

    ''' Log items '''
    train_loss = AverageValueMeter()
    val_loss = AverageValueMeter()
    with open(logname, 'a') as f: #open and append
        f.write(str(model) + '\n')

    train_curve = []
    val_curve = []
    labels_generated_points = torch.Tensor(range(1, (opt.n_primitives+1)*(opt.out_num_points//opt.n_primitives)+1)).view(opt.out_num_points//opt.n_primitives,(opt.n_primitives+1)).transpose(0,1)
    labels_generated_points = (labels_generated_points)%(opt.n_primitives+1)
    labels_generated_points = labels_generated_points.contiguous().view(-1)

    ''' Train '''
    if not opt.eval:
        for epoch in range(opt.nepoch):
            #TRAIN MODE
            train_loss.reset()

            # learning rate schedule
            if epoch == 20:
                optimizer = optim.Adam(model.parameters(), lr=lrate/10.0)
            if epoch == 40:
                optimizer = optim.Adam(model.parameters(), lr=lrate/100.0)

            emd, chamfer = train_one_epoch(network,
                                           optimizer,
                                           dataloader,
                                           train_loss,
                                           vis)
            print('Train [{}] EMD:{}, CD:{}'.format(epoch, emd, chamfer))
            train_curve.append(train_loss.avg)

            # Validation
            if epoch % 5 == 0:
                val_loss.reset()
                emd, chamfer = evaluate(network, dataloader_test, val_loss, vis)
                print('Val [{}] EMD: {} CD: {}'.format(epoch, emd, chamfer))
            val_curve.append(val_loss.avg)

            # Epoch logging
            vis.line(X=np.column_stack((np.arange(len(train_curve)),np.arange(len(val_curve)))),
                        Y=np.column_stack((np.array(train_curve),np.array(val_curve))),
                        win='loss',
                        opts=dict(title="emd", legend=["train_curve" + opt.exp_name, "val_curve" + opt.exp_name], markersize=2, ), )
            vis.line(X=np.column_stack((np.arange(len(train_curve)),np.arange(len(val_curve)))),
                        Y=np.log(np.column_stack((np.array(train_curve),np.array(val_curve)))),
                        win='log',
                        opts=dict(title="log_emd", legend=["train_curve"+ opt.exp_name, "val_curve"+ opt.exp_name], markersize=2, ), )

            log_table = {
                "train_loss" : train_loss.avg,
                "val_loss" : val_loss.avg,
                "epoch" : epoch,
                "lr" : lrate,
            }
            with open(logname, 'a') as f:
                f.write('json_stats: ' + json.dumps(log_table) + '\n')

            # Save network
            torch.save(model.state_dict(), '%s/network.pth' % (log_dir))
            print('Model saved.')
    else:
        emd, chamfer = evaluate(network, dataloader_test, val_loss, vis,
                                precise=True)
        print('Eval EMD: {} CD: {}'.format(emd, chamfer))

max_stars_repo_name: itsumu/MSN-Point-Cloud-Completion

max_issues_repo_name: itsumu/MSN-Point-Cloud-Completion

max_forks_repo_name: itsumu/MSN-Point-Cloud-Completion


text:
import numpy as np
from mpl_toolkits.basemap import Basemap
from numpy import ma
import pygrib, sys
import matplotlib.pyplot as plt

for grb in pygrib.open('../sampledata/tigge.grb'):
    fld = 0.01*grb['values'] # convert to hPa
    lats,lons = grb.latlons()
    sys.stdout.write('%s %s %s %s' % \
            (grb['centre'], fld.shape, fld.min(), fld.max()))
    fig=plt.figure(figsize=(10,5))
    fig.add_axes([0.1,0.1,0.8,0.8])
    m = Basemap(projection='cyl',lon_0=180)
    x, y = m(lons,lats)
    levels = np.arange(475,1101,25)
    CS = m.contourf(x,y,fld,levels,cmap=plt.cm.jet)
    plt.colorbar(drawedges=True, shrink=0.8) # draw colorbar
    m.drawcoastlines()
    m.drawparallels(np.arange(-80,81,20),labels=[1,0,0,0])
    m.drawmeridians(np.arange(0,360,60),labels=[0,0,0,1])
    m.drawmapboundary()
    plt.title(grb['name']+': '+grb['centre'].upper(),fontsize=12)
plt.show()

reason: import numpy,from numpy