UDP Get Image

getImageUdp.py

# Written by Brian Ejike (2017)
# Distributed under the MIT License

import serial, time, socket

WIDTH = 256
HEIGHT = 288
READ_LEN = int(WIDTH * HEIGHT / 2)

DEPTH = 8
HEADER_SZ = 54

portSettings = ["", 0]


print("----------Extract Fingerprint Image------------")
print()


# assemble bmp header for a grayscale image
def assembleHeader(width, height, depth, cTable=False):
    header = bytearray(HEADER_SZ)
    header[0:2] = b"BM"  # bmp signature
    byte_width = int((depth * width + 31) / 32) * 4
    if cTable:
        header[2:6] = ((byte_width * height) + (2**depth) * 4 + HEADER_SZ).to_bytes(
            4, byteorder="little"
        )  # file size
    else:
        header[2:6] = ((byte_width * height) + HEADER_SZ).to_bytes(
            4, byteorder="little"
        )  # file size
    # header[6:10] = (0).to_bytes(4, byteorder='little')
    if cTable:
        header[10:14] = ((2**depth) * 4 + HEADER_SZ).to_bytes(
            4, byteorder="little"
        )  # offset
    else:
        header[10:14] = (HEADER_SZ).to_bytes(4, byteorder="little")  # offset

    header[14:18] = (40).to_bytes(4, byteorder="little")  # file header size
    header[18:22] = width.to_bytes(4, byteorder="little")  # width
    header[22:26] = (-height).to_bytes(4, byteorder="little", signed=True)  # height
    header[26:28] = (1).to_bytes(2, byteorder="little")  # no of planes
    header[28:30] = depth.to_bytes(2, byteorder="little")  # depth
    # header[30:34] = (0).to_bytes(4, byteorder='little')
    header[34:38] = (byte_width * height).to_bytes(4, byteorder="little")  # image size
    header[38:42] = (1).to_bytes(4, byteorder="little")  # resolution
    header[42:46] = (1).to_bytes(4, byteorder="little")
    # header[46:50] = (0).to_bytes(4, byteorder='little')
    # header[50:54] = (0).to_bytes(4, byteorder='little')
    return header


def options():
    print("Options:")
    print("\tPress 1 to enter serial port settings")
    print("\tPress 2 to scan a fingerprint and save the image")
    print("\tPress 3 to view help")
    print("\tPress 4 to exit")
    print()
    choice = input(">> ")
    print()
    return choice


# def getSettings():
#     portSettings[0] = input("Enter Arduino serial port number: ")
#     portSettings[1] = int(input("Enter serial port baud rate: "))
#     print()


def getPrint():
    """
    First enter the port settings with menu option 1:
    >>> Enter Arduino serial port number: COM13
    >>> Enter serial port baud rate: 57600

    Then enter the filename of the image with menu option 2:
    >>> Enter filename/path of output file (without extension): myprints
    Found fingerprint sensor!
    .
    .
    .
    (Here you communicate with the Arduino and follow instructions)
    .
    .
    .
    Extracting image...saved as <filename>.bmp

    """
    out = open(
        input("Enter filename/path of output file (without extension): ") + ".bmp", "wb"
    )
    # assemble and write the BMP header to the file
    out.write(assembleHeader(WIDTH, HEIGHT, DEPTH, True))
    try:
        sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        sock.bind(("192.168.1.9", 9999))
    except socket.error:
        print("Error socket")
    for i in range(256):
        # write the colour palette
        out.write(i.to_bytes(1, byteorder="little") * 4)
        
    while True:
        try:
            written = 0
            last_read = time.time()
            
            while written < READ_LEN:  # start recving image
                # read a packet of arbitrary length from the socket; supposedly blocking call.
                data, addr = sock.recvfrom(4096)
                
                # if we get nothing more after a 5 sec timeout
                if not data:
                    if written and time.time() - last_read > 5:
                        print("Timeout!")
                        out.close()  # close port and file
                        sock.close()
                        return False
                    else:
                        continue
                    
                # since each byte holds colour info for 2 adjacent pixels in each nibble,
                # assume that both pixels are nearly the same colour and so repeat each byte twice to get 2 pixels data
                
                image_data = [byte for byte in data for i in range(2)]
                image_data = bytes(image_data)
                out.write(image_data)
                
                written += len(data)

            out.close()
            sock.close()  # close file
            print("Image saved as ", out.name, ", received size = ", written, "bytes.")
            return True
        except Exception as e:
            print("Read failed: ", e)
            out.close()
            sock.close()
            return False
        except KeyboardInterrupt:
            print("Closing port.")
            out.close()
            sock.close()
            return False


while True:
    chose = options()
    if chose == "4":
        break
    # elif chose == "1":
    #     getSettings()
    elif chose == "2":
        res = getPrint()
        if not res:
            print("Image extraction failed!")
        continue
    elif chose == "3":
        print("================= HELP ==================")
        print(getPrint.__doc__)
        print("=========================================")