def select_rgb_white_yellow(image):
    # 過濾背景
    lower = np.uint8([120, 120, 120])
    upper = np.uint8([255, 255, 255])
    # 三個(gè)通道內(nèi)，低于lower和高于upper的部分分別變成0， 在lower-upper之間的值變成255， 相當(dāng)于mask，過濾背景
    # 保留了像素值在120-255之間的像素值
    white_mask = cv2.inRange(image, lower, upper)
    masked_img = cv2.bitwise_and(image, image, mask=white_mask)
    return masked_img
masked_img = select_rgb_white_yellow(test_image)

low_threshold, high_threshold = 50, 200
edges_img = cv2.Canny(gray_img, low_threshold, high_threshold)

def select_region(image):
    """這里手動選擇區(qū)域"""
    rows, cols = image.shape[:2]
    
    # 下面定義6個(gè)標(biāo)定點(diǎn), 這個(gè)點(diǎn)的順序必須讓它化成一個(gè)區(qū)域，如果調(diào)整，可能會交叉起來，所以不要動
    pt_1  = [cols*0.06, rows*0.90]   # 左下
    pt_2 = [cols*0.06, rows*0.70]    # 左上
    pt_3 = [cols*0.32, rows*0.51]    # 中左
    pt_4 = [cols*0.6, rows*0.1]      # 中右
    pt_5 = [cols*0.90, rows*0.1]     # 右上
    pt_6 = [cols*0.90, rows*0.90]    # 右下
    
    vertices = np.array([[pt_1, pt_2, pt_3, pt_4, pt_5, pt_6]], dtype=np.int32)
    point_img = image.copy()
    point_img = cv2.cvtColor(point_img, cv2.COLOR_GRAY2BGR)
    for point in vertices[0]:
        cv2.circle(point_img, (point[0], point[1]), 10, (0, 0, 255), 4)
    # cv_imshow('points_img', point_img)
    
    # 定義mask矩陣， 只保留點(diǎn)內(nèi)部的區(qū)域
    mask = np.zeros_like(image)
    if len(mask.shape) == 2:
        cv2.fillPoly(mask, vertices, 255)   # 點(diǎn)框住的地方填充為白色
        #cv_imshow('mask', mask)
    roi_image = cv2.bitwise_and(image, mask)
    return roi_image

roi_image = select_region(edges_img)

def hough_lines(image):
    # 輸入的圖像需要是邊緣檢測后的結(jié)果
    # minLineLengh(線的最短長度，比這個(gè)短的都被忽略)和MaxLineCap（兩條直線之間的最大間隔，小于此值，認(rèn)為是一條直線）
    # rho距離精度,theta角度精度,threshod超過設(shè)定閾值才被檢測出線段
    return cv2.HoughLinesP(image, rho=0.1, theta=np.pi/10, threshold=15, minLineLength=9, maxLineGap=4)

list_of_lines = hough_lines(roi_image)  # (2338, 1, 4)

def identity_blocks(image, lines, make_copy=True):
    if make_copy:
        new_image = image.copy()
    
    # 過濾部分直線
    stayed_lines = []
    for line in lines:
        for x1, y1, x2, y2 in line:
            # 這里是過濾直線，必須保證不能是斜的線，且水平方向不能太長或者太短
            if abs(y2-y1) <=1 and abs(x2-x1) >=25 and abs(x2-x1) <= 55:
                stayed_lines.append((x1,y1,x2,y2))
    
    # 對直線按照x1排序, 這樣能讓這些線從上到下排列好， 這個(gè)排序是從第一列的第一條橫線，往下走，然后是第二列第一條橫線往下，...
    list1 = sorted(stayed_lines, key=operator.itemgetter(0, 1))

	# 找到多個(gè)列，相當(dāng)于每列是一排車
    clusters = collections.defaultdict(list)
    dIndex = 0
    clus_dist = 10   # 每一列之間的那個(gè)距離
    for i in range(len(list1) - 1):
        # 看看相鄰兩條線之間的距離，如果是一列的，那么x1這個(gè)距離應(yīng)該很近，畢竟是同一列上的
        # 如果這個(gè)值大于10了，說明是下一列的了，此時(shí)需要移動dIndex， 這個(gè)表示的是第幾列 
        distance = abs(list1[i+1][0] - list1[i][0])
        if distance <= clus_dist:
            clusters[dIndex].append(list1[i])
            clusters[dIndex].append(list1[i+1])
        else:
            dIndex += 1

	# 得到每列停車位的矩形框
    rects = {}
    i = 0
    for key in clusters:
        all_list = clusters[key]
        cleaned = list(set(all_list))
        # 有5個(gè)停車位至少
        if len(cleaned) > 5:
            cleaned = sorted(cleaned, key=lambda tup: tup[1])
            avg_y1 = cleaned[0][1]
            avg_y2 = cleaned[-1][1]
            if abs(avg_y2-avg_y1) < 15:
                continue
            avg_x1 = 0
            avg_x2 = 0
            for tup in cleaned:
                avg_x1 += tup[0]
                avg_x2 += tup[2]
            avg_x1 = avg_x1 / len(cleaned)
            avg_x2 = avg_x2 / len(cleaned)
            
            rects[i] = [avg_x1, avg_y1, avg_x2, avg_y2]
            i += 1
    print('Num Parking Lanes: ', len(rects))

	# 把列矩形畫出來
    buff = 7
    for key in rects:
        tup_topLeft = (int(rects[key][0] - buff), int(rects[key][1]))
        tup_botRight = (int(rects[key][2] + buff), int(rects[key][3]))
        cv2.rectangle(new_image, tup_topLeft, tup_botRight, (0, 255, 0), 3)
    return new_image, rects

def rect_finetune(image, rects, copy_img=True):
    if copy_img:
        image_copy = image.copy()
    # 下面需要對上面的框進(jìn)行坐標(biāo)微調(diào)， 讓框更加準(zhǔn)確
    # 這個(gè)框很重要，影響后面停車位的統(tǒng)計(jì)，盡量不能有遺漏
    for k in rects:
        if k == 0:
            rects[k][1] -= 10
        elif k == 1:
            rects[k][1] -= 10
            rects[k][3] -= 10
        elif k == 2 or k == 3 or k == 5:
            rects[k][1] -= 4
            rects[k][3] += 13
        elif k == 6 or k == 8:
            rects[k][1] -= 18
            rects[k][3] += 12
        elif k == 9:
            rects[k][1] += 10
            rects[k][3] += 10
        elif k == 10:
            rects[k][1] += 45
        elif k == 11:
            rects[k][3] += 45
    
    buff = 8
    for key in rects:
        tup_topLeft = (int(rects[key][0]-buff), int(rects[key][1]))
        tup_botRight = (int(rects[key][2]+buff), int(rects[key][3]))
        cv2.rectangle(image_copy, tup_topLeft, tup_botRight, (0, 255, 0), 3)
    
    return image_copy, rects

def draw_parking(image, rects, make_copy=True, save=True):
    gap = 15.5
    spot_dict = {}  # 一個(gè)車位對應(yīng)一個(gè)位置
    tot_spots = 0
    
    #微調(diào)
    adj_x1 = {0: -8, 1:-15, 2:-15, 3:-15, 4:-15, 5:-15, 6:-15, 7:-15, 8:-10, 9:-10, 10:-10, 11:0}
    adj_x2 = {0: 0, 1: 15, 2:15, 3:15, 4:15, 5:15, 6:15, 7:15, 8:10, 9:10, 10:10, 11:0}
    fine_tune_y = {0: 4, 1: -2, 2: 3, 3: 1, 4: -3, 5: 1, 6: 5, 7: -3, 8: 0, 9: 5, 10: 4, 11: 0}
    
    for key in rects:
        tup = rects[key]
        x1 = int(tup[0] + adj_x1[key])
        x2 = int(tup[2] + adj_x2[key])
        y1 = int(tup[1])
        y2 = int(tup[3])
        cv2.rectangle(new_image, (x1, y1),(x2,y2),(0,255,0),2)
        
        num_splits = int(abs(y2-y1)//gap)
        for i in range(0, num_splits+1):
            y = int(y1+i*gap) + fine_tune_y[key]
            cv2.line(new_image, (x1, y), (x2, y), (255, 0, 0), 2)
        if key > 0 and key < len(rects) - 1:
            # 豎直線
            x = int((x1+x2) / 2)
            cv2.line(new_image, (x, y), (x, y2), (0, 0, 255), 2)
        
        # 計(jì)算數(shù)量   除了第一列和最后一列，中間的都是兩列的
        if key == 0 or key == len(rects) - 1:
            tot_spots += num_splits + 1
        else:
            tot_spots += 2 * (num_splits + 1)
        
        # 字典對應(yīng)好
        if key == 0 or key == len(rects) - 1:
            for i in range(0, num_splits+1):
                cur_len = len(spot_dict)
                y = int(y1 + i * gap) + fine_tune_y[key]
                spot_dict[(x1, y, x2, y+gap)] = cur_len + 1
        else:
            for i in range(0, num_splits+1):
                cur_len = len(spot_dict)
                y = int(y1 + i * gap) + fine_tune_y[key]
                x = int((x1+x2) / 2)
                spot_dict[(x1, y, x, y+gap)] = cur_len + 1
                spot_dict[(x, y, x2, y+gap)] = cur_len + 2
  
    return new_image, spot_dict

# 去掉多余的停車位
invalid_spots = [10, 11, 33, 34, 37, 38, 61, 62, 93, 94, 95, 97, 98, 135, 137, 138, 187, 249, 
           250, 253, 254, 323, 324, 327, 328, 467, 468, 531, 532]
valid_spots_dict = {}
cur_idx = 1
for k, v in spot_dict.items():
    if v in invalid_spots:
        continue
    valid_spots_dict[k] = cur_idx
    cur_idx += 1

# 把每一個(gè)有效停車位標(biāo)記出來
tmp_img = test_image.copy()
for k, v in valid_spots_dict.items():
    cv2.rectangle(tmp_img, (int(k[0]), int(k[1])),(int(k[2]),int(k[3])), (0,255,0) , 2)
cv_imshow('valid_pot', tmp_img)

def save_images_for_cnn(image, spot_dict, folder_name = '../cnn_pred_data'):
    for spot in spot_dict.keys():
        (x1, y1, x2, y2) = spot
        (x1, y1, x2, y2) = (int(x1), int(y1), int(x2), int(y2))
        
        # 裁剪
        spot_img = image[y1:y2, x1:x2]
        spot_img = cv2.resize(spot_img, (0, 0), fx=2.0, fy=2.0)
        spot_id = spot_dict[spot]
        
        filename = 'spot_{}.jpg'.format(str(spot_id))
        
        # print(spot_img.shape, filename, (x1,x2,y1,y2))
        cv2.imwrite(os.path.join(folder_name, filename), spot_img)
  
save_images_for_cnn(test_image, valid_spots_dict)

def train_model():
    # 獲取dataloader
    data_root = os.getcwd()
    image_path = os.path.join(data_root, "train_data")
    train_data_path = os.path.join(image_path, "train")
    val_data_path = os.path.join(image_path, "test")
    train_loader, validat_loader, train_num, val_num = get_dataloader(train_data_path, val_data_path,
                                                                      data_transform_pretrain, batch_size=8)

    # 創(chuàng)建模型 注意這里沒指定類的個(gè)數(shù)，默認(rèn)是1000類
    net = resnet34()
    model_weight_path = 'saved_model_weight/resnet34_pretrain_ori_low_torch_version.pth'

    # 使用預(yù)訓(xùn)練的參數(shù)，然后進(jìn)行finetune
    net.load_state_dict(torch.load(model_weight_path, map_location='cpu'))

    # 改變fc layer structure  把fc的輸出維度改為2
    in_channel = net.fc.in_features
    net.fc = nn.Linear(in_channel, 2)
    net.to(device)

    # 模型訓(xùn)練配置
    loss_function = nn.CrossEntropyLoss()
    optimizer = optim.Adam(net.parameters(), lr=0.0001)

    epochs = 30
    save_path = "saved_model_weight/resnet34_pretrain.pth"
    best_acc = 0.
    train_steps = len(train_loader)

    model_train(net, train_loader, validat_loader, epochs, device, optimizer, loss_function, train_steps, val_num,
                save_path, best_acc)

data_transform_pretrain = {
        "train": transforms.Compose([
            transforms.RandomResizedCrop(224),  # 對圖像隨機(jī)裁剪， 訓(xùn)練集用，驗(yàn)證集不用
            transforms.RandomHorizontalFlip(),
            transforms.ToTensor(),
            # 這里的中心化處理參數(shù)需要官方給定的參數(shù)，這里是ImageNet圖片的各個(gè)通道的均值和方差，不能隨意指定了
            transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225))
        ]),
        "val": transforms.Compose([
            # 驗(yàn)證過程中，這里也進(jìn)行了一點(diǎn)點(diǎn)改動
            transforms.Resize(256),
            transforms.CenterCrop(224),
            transforms.ToTensor(),
            transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225))
        ]),
        "test": transforms.Compose([
            transforms.Resize(256),
            transforms.CenterCrop(224),
            transforms.ToTensor(),
            transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225))
        ])
    }

xxx.pt is a zip archive(did you mean to use torch.jit.load()?)“

model_weight_path = "saved_models/resnet34_pretrain_ori.pth"
    state_dict = torch.load(model_weight_path)
    torch.save(state_dict, 'saved_models/resnet34_pretrain_ori_low_torch_version.pth', _use_new_zipfile_serialization=False)

# 創(chuàng)建模型 注意這里沒指定類的個(gè)數(shù)，默認(rèn)是1000類
net = resnet34()
model_weight_path = 'saved_model_weight/resnet34_pretrain_ori_low_torch_version.pth'

# 使用預(yù)訓(xùn)練的參數(shù)，然后進(jìn)行finetune
net.load_state_dict(torch.load(model_weight_path, map_location='cpu'))

# 改變fc layer structure  把fc的輸出維度改為2
in_channel = net.fc.in_features
net.fc = nn.Linear(in_channel, 2)
net.to(device)

# 模型訓(xùn)練配置
loss_function = nn.CrossEntropyLoss()
# 訓(xùn)練的時(shí)候，也可以凍結(jié)掉卷積層的參數(shù), 也可以指定不同層的參數(shù)使用不同的學(xué)習(xí)率進(jìn)行訓(xùn)練
res_params, conv_params, fc_params = [], [], []
# named_parameters()能返回每一層的名字以及參數(shù)，是一個(gè)字典
for name, param in net.named_parameters():
     # layer 系列是殘差層
     if ('layer' in name):
           res_params.append(param)
     # 全連接層
     elif ('fc' in name):
           fc_params.append(param)
     else:
           param.requires_grad = False

params = [
     {'params': res_params, 'lr': 0.0001},
     {'params': fc_params, 'lr': 0.0002},
]

optimizer = optim.Adam(params)

def predict_on_img(img, spot_dict, model, class_indict, make_copy=True, color=[0, 255, 0], alpha=0.5, save=True):
    # 這個(gè)是停車場的全景圖像
    if make_copy:
        new_image = np.copy(img)
        overlay = np.copy(img)

    cnt_empty, all_spots = 0, 0
    for spot in tqdm(spot_dict.keys()):
        all_spots += 1
        (x1, y1, x2, y2) = spot
        (x1, y1, x2, y2) = (int(x1), int(y1), int(x2), int(y2))
        spot_img = img[y1:y2, x1:x2]
        spot_img_pil = Image.fromarray(spot_img)

        label = model_infer(spot_img_pil, model, class_indict)
        if label == 'empty':
            cv2.rectangle(overlay, (int(x1), int(y1)), (int(x2), int(y2)), color, -1)
            cnt_empty += 1

    cv2.addWeighted(overlay, alpha, new_image, 1 - alpha, 0, new_image)

    # 顯示結(jié)果的
    cv2.putText(new_image, "Available: %d spots" % cnt_empty, (30, 95),
                cv2.FONT_HERSHEY_SIMPLEX,
                0.7, (255, 255, 255), 2)
    cv2.putText(new_image, "Total: %d spots" % all_spots, (30, 125),
                cv2.FONT_HERSHEY_SIMPLEX,
                0.7, (255, 255, 255), 2)

    if save:
        filename = 'with_marking_predict.jpg'
        cv2.imwrite(filename, new_image)
    # cv_imshow('new_image', new_image)
    return new_image

def predict_on_video(video_path, spot_dict, model, class_indict, ret=True):
    cap = cv2.VideoCapture(video_path)
    count = 0
    while ret:
        ret, image = cap.read()
        count += 1

        if count == 5:
            count = 0
            new_image = predict_on_img(image, spot_dict, model, class_indict, save=False)

            cv2.imshow('frame', new_image)
            if cv2.waitKey(10) & 0xFF == ord('q'):
                break

    cv2.destroyAllWindows()
    cap.release()