import time
import cv2
import os
import base64
from socket import *


class Client:
    def __init__(self, host="127.0.0.1", port=8001):
        self.host = host
        self.port = port
        self.socket = self.connect()

    def connect(self):
        tcp_socket = socket()
        tcp_socket.connect((self.host, self.port))
        return tcp_socket
    
    # 發(fā)送的是編碼，有些項目中就不需要保存路徑，直接發(fā)送編碼了
    def base64_encode_img(self, img_data):
        _, image_encoded = cv2.imencode('.jpg', img_data)
        image_bytes = image_encoded.tobytes()
        image_base64 = base64.b64encode(image_bytes).decode('utf-8')
        return image_base64

    def get_img(self, img_path):
        img = cv2.imread(img_path)
        img_d = self.base64_encode_img(img)
        img_b = img_d.encode()
        return img_b

    def send_msg(self, msg):
        self.socket.send(msg.encode())

    def send_img(self, img_path):
        img_data = self.get_img(img_path)
        n = 0
        while True:
            self.socket.send(img_data[n:n + 1024])
            n += 1024
            if n > len(img_data):
                print(f"已發(fā)送:{img_path}")
                break

    def close(self):
        self.socket.close()


if __name__ == '__main__':

    path = "C:/Users/cheng/Desktop/Spider/dilireba"   # 圖片路徑
    # path = "C:/Users/cheng/Desktop/Spider/gulinazha"
    # path = "C:/Users/cheng/Desktop/Spider/boduo"
    pics = os.listdir(path)
    for pic in pics:
        C = Client(host="127.0.0.1", port=8001)
        p_p = path + "/" + pic
        C.send_img(p_p)  # 發(fā)送單張圖片
        C.close()

import os
from socket import *
import base64


# 圖片解碼
def base64_decode_img(base64_code):
    """
    :param base64_code: base64編碼
    :return: 二進制據(jù),直接寫入保存為圖片
    """
    img_bytes = base64.b64decode(base64_code)
    return img_bytes


class WebSocketServer:
    def __init__(self, host="127.0.0.1", port=8001):
        self.socket = socket()
        self.host = host
        self.port = port

    def bind(self):
        """監(jiān)聽套接字"""
        self.socket.bind((self.host, self.port))
        self.socket.listen()  # 設(shè)置監(jiān)聽套接字

    def bind_connection(self):
        """建立連接"""
        while True:
            print("等待連接....")
            conn, addr = self.socket.accept()  # 連接 conn是連接對象；addr=(host,port)
            print(f"連接到:{addr[0]}\t端口:{addr[1]}")
            return conn, addr

    def close(self):
        """ 關(guān)閉監(jiān)聽套接字"""
        self.socket.close()


class Connection:
    def __init__(self, conn, addr):
        self.conn = conn
        self.addr = addr
        self.save_path = "recv/"
        self.flag = False

    def close(self):
        """斷開連接"""
        self.conn.close()
        print(f"{self.addr[0]}:{self.addr[1]}斷開連接")

    def send_msg(self, message):
        """發(fā)送信息"""
        self.conn.send(message.encode())

    def receive_msg(self, bufsize=1024):
        """接受信息"""
        data_received = self.conn.recv(bufsize).decode()
        return data_received

    def handle(self, count):
        if not os.path.exists(self.save_path):
            os.makedirs(self.save_path)
        img_data = "" # 需要存儲的圖片數(shù)據(jù)
        while True:
            data = self.receive_msg()  # 持續(xù)獲取該對象信息
            if data:
                img_data += data
            else:
                self.flag = True
                break
        d = base64_decode_img(img_data)
        with open(f"{self.save_path}/{count}.jpg", "wb") as f:
            f.write(d) if d else None


if __name__ == '__main__':
    webserver = WebSocketServer(host="127.0.0.1", port=8001)
    webserver.bind()
    count = 0
    while True:
        conn, addr = webserver.bind_connection()
        con = Connection(conn, addr)
        while True:
            con.handle(count)
            count += 1
            print(f"完成{count}")
            if con.flag: break