Receiving a full image from a TCP socket communication

I have a TCP communication where my client continuously sends images as a byte array to a server and receives a response back, my problem is that when the server receives the images they are not done being received even though I've added a flag to indicate the end of the image. I'd like to know a better way to ensure that the image file is received completely before receiving a new one

EDIT: My new attempt:

Client.py

import numpy as np
import cv2
from PIL import Image
import base64
import socket
def main(data):
    s = socket.socket()
    s.connect(("127.0.0.1", 999))

    decoded_data = base64.b64decode(data)

    print("Sending...")
    s.sendall(decoded_data)
    s.shutdown(s.SHUTWR)

    b_data = b''
    while True:
        txt_data = s.recv(2048)
        if not txt_data: break
        b_data += txt_data
        print('response received from the server: ' + b_data.decode())
    return b_data.decode()

Server.py

import socket
from PIL import Image
import io
import numpy as np
import cv2
import uuid

IP = '127.0.0.1'
PORT = 999

with socket.socket() as s:
    s.bind((IP,PORT))
    s.listen(1)
    count = 0
    print ('The server is ready')
    while True:
        con, addr = s.accept()
        filename = str(uuid.uuid4())
        count = count + 1
        img_dir = 'C:/Users/my_user/stream_images/'
        img_format = '.png'
        with con:
            img = b''
            while True:
                data = con.recv(2048)
                if not data:
                    break
                img += data
            image_name = img_dir+'frame'+str(count)+img_format
            pil_image = io.BytesIO(img)
            img = np.array(Image.open(pil_image))
            img = cv2.rotate(img, cv2.ROTATE_90_CLOCKWISE)
            cCode = str('Thank you for connecting')
            con.sendall(cCode.encode())
            print("called con.sendall")
            cv2.imshow('frame', img)
            if cv2.waitKey(1) & 0xFF == ord('q'):
                break
cv2.destroyAllWindows()

Currently, I am now able to fully send the images and receive them properly at the server, the only problem is that I am no longer sending a response back after the image is received, so there is something wrong with how I am receiving the reply message at the client side.

Solution

As user207421 suggested you can shutdown the socket for writing after sending the image on the client-side while still being able to receive an confirmatory answer from the server. Another problem you're facing here is the blocking nature of cv2.waitKey, which essentially halts the server until the user presses q in the cv2 window (the server will not be able to handle any other requests). I'd generally recommend to separate your network/IO logic from user interface logic. To circumvent the blocking behaviour of I've implemented a very basic image_viewer, which waits for incoming images in a thread that runs separately from the server loop by passing images through a Queue.

The client code looks as follows:

import socket
from PIL import Image

def send_image(img: Image, host: str = '127.0.0.1', port: int = 999):
    s = socket.socket()
    s.connect((host, port))
    img_data = img._repr_png_()
    s.sendall(img_data)
    s.shutdown(socket.SHUT_WR)  # close socket for writing, receiving is still possible
    print(f'Sent {len(img_data) / 1024:,.1f} kB of image data.')
    b_data = b''
    while recv_data := s.recv(2048):
        b_data += recv_data
    print(f'Server response: {b_data.decode()}')
    # maybe check server response for server side errors etc. and add return value for this function?

# use like: send_image(Image.open('test.png'))

The server code is:

import io
import queue
import socket
import threading

import cv2
import numpy as np
from PIL import Image

def image_viewer(q: queue.Queue):
    while True:
        try:
            img_name, img = q.get(block=True, timeout=.1)  # poll every 0.1 seconds
            print(f'Image viewer: displaying `{img_name}`!')
            cv2.imshow('Image preview', img)
        except queue.Empty:
            ...  # no new image to display
        key = cv2.pollKey()  # non-blocking
        if key & 0xff == ord('q'):
            cv2.destroyAllWindows()
            print('Image viewer was closed')
            return

def serve_forever(host: str, port: int, img_dir: str = 'C:/Users/my_user/stream_images/', img_format: str = '.png'):
    q = queue.Queue()
    img_viewer = threading.Thread(target=image_viewer, args=(q,))
    img_viewer.start()

    with socket.socket() as s:
        s.bind((host, port))
        s.listen(1)
        count = 0
        print('The server is ready')
        while True:
            conn, addr = s.accept()
            count = count + 1
            img_name = img_dir + 'frame' + str(count) + img_format
            print (f'Client connected: {addr}')
            img = b''
            while data := conn.recv(2048):
                img += data
            conn.sendall('Thank you for connecting'.encode())  # maybe use return codes for success, error etc.?
            conn.close()
            pil_img = Image.open(io.BytesIO(img))  # might want to save to disk?
            np_img = np.asarray(pil_img)
            np_img = cv2.rotate(np_img, cv2.ROTATE_90_CLOCKWISE)
            q.put((img_name, np_img))
            print (f'Client at {addr} disconnected after receiving {len(img) / 1024:,.1f} kB of data.')

if __name__ == '__main__':
    serve_forever('127.0.0.1', 999)