Python中Numpy三维ndarray的传输压缩方法求助

Great question! When dealing with 3D NumPy arrays that have lots of spatial redundancy (similar neighboring regions, like you described), there are several effective compression strategies to cut down on data size during socket transfer. Below are the most practical options, with actionable code snippets to implement them:

1. Zlib/Gzip (Built-in, General-Purpose) #

This is the easiest starting point since Python has built-in zlib and gzip modules—no extra installs needed. It works well for arrays with redundant data, balancing compression ratio and speed.

Sender code snippet:

import numpy as np
import zlib
import socket

# Replace with your actual 3D array
arr = np.random.rand(100, 100, 100).astype(np.float32)

# Convert array to raw bytes and compress
arr_bytes = arr.tobytes()
compressed_data = zlib.compress(arr_bytes, level=zlib.Z_BEST_COMPRESSION)

# Send over socket: first transmit the length of compressed data, then the data itself
sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
sock.connect(("receiver_ip", 12345))
# Send length as a 4-byte big-endian integer
sock.sendall(len(compressed_data).to_bytes(4, byteorder='big'))
sock.sendall(compressed_data)
sock.close()

Receiver code snippet:

import numpy as np
import zlib
import socket

sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
sock.bind(("0.0.0.0", 12345))
sock.listen(1)
conn, addr = sock.accept()

# First read the length of compressed data
len_data = conn.recv(4)
compressed_length = int.from_bytes(len_data, byteorder='big')

# Receive the full compressed data in chunks
compressed_data = b''
while len(compressed_data) < compressed_length:
    chunk = conn.recv(min(4096, compressed_length - len(compressed_data)))
    if not chunk:
        break
    compressed_data += chunk

# Decompress and restore the original array
arr_bytes = zlib.decompress(compressed_data)
arr = np.frombuffer(arr_bytes, dtype=np.float32).reshape(100, 100, 100)
conn.close()
sock.close()

2. LZ4/Snappy (High-Speed Compression) #

If you prioritize transfer speed over maximum compression ratio, LZ4 or Snappy are excellent choices—they’re much faster than zlib while still providing meaningful compression for redundant data.

First install the library:
pip install lz4 (for LZ4) or pip install python-snappy (for Snappy)

Sender example with LZ4:

import numpy as np
import lz4.frame
import socket

arr = np.random.rand(100, 100, 100).astype(np.float32)
arr_bytes = arr.tobytes()
compressed_data = lz4.frame.compress(arr_bytes, compression_level=lz4.frame.COMPRESSIONLEVEL_MAX)

# Send length + data (same socket logic as zlib example)
sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
sock.connect(("receiver_ip", 12345))
sock.sendall(len(compressed_data).to_bytes(4, byteorder='big'))
sock.sendall(compressed_data)
sock.close()

Receiver example with LZ4:

import numpy as np
import lz4.frame
import socket

sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
sock.bind(("0.0.0.0", 12345))
sock.listen(1)
conn, addr = sock.accept()

len_data = conn.recv(4)
compressed_length = int.from_bytes(len_data, byteorder='big')
compressed_data = b''
while len(compressed_data) < compressed_length:
    chunk = conn.recv(min(4096, compressed_length - len(compressed_data)))
    if not chunk:
        break
    compressed_data += chunk

arr_bytes = lz4.frame.decompress(compressed_data)
arr = np.frombuffer(arr_bytes, dtype=np.float32).reshape(100, 100, 100)
conn.close()
sock.close()

3. HDF5 with Compression (Structured Data) #

If your 3D array has metadata (like axis labels, dtype info) or you need to handle very large arrays, using HDF5 (via h5py) is a robust option. It supports built-in compression and preserves all array attributes automatically.

Install first: pip install h5py

Sender code:

import numpy as np
import h5py
import io
import socket

arr = np.random.rand(100, 100, 100).astype(np.float32)

# Write array to an in-memory HDF5 file with gzip compression
buffer = io.BytesIO()
with h5py.File(buffer, 'w') as f:
    f.create_dataset('data', data=arr, compression='gzip', compression_opts=9)
buffer.seek(0)
compressed_data = buffer.read()

# Send over socket
sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
sock.connect(("receiver_ip", 12345))
sock.sendall(len(compressed_data).to_bytes(4, byteorder='big'))
sock.sendall(compressed_data)
sock.close()

Receiver code:

import numpy as np
import h5py
import io
import socket

sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
sock.bind(("0.0.0.0", 12345))
sock.listen(1)
conn, addr = sock.accept()

len_data = conn.recv(4)
compressed_length = int.from_bytes(len_data, byteorder='big')
compressed_data = b''
while len(compressed_data) < compressed_length:
    chunk = conn.recv(min(4096, compressed_length - len(compressed_data)))
    if not chunk:
        break
    compressed_data += chunk

# Read from in-memory HDF5 buffer
buffer = io.BytesIO(compressed_data)
with h5py.File(buffer, 'r') as f:
    arr = f['data'][()]
conn.close()
sock.close()

4. Image-Specific Compression (For Volumetric Image Data) #

If your 3D array is volumetric image data (e.g., (frames, height, width, channels)), you can leverage 2D image compression algorithms like JPEG/PNG by processing each slice. This works great for data where each 2D plane has strong spatial redundancy.

Using Pillow: pip install pillow

Sender code:

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

# Sample 3D image array: (100 frames, 100x100, 3 RGB channels)
arr = np.random.randint(0, 255, size=(100, 100, 100, 3), dtype=np.uint8)

# Compress each frame as JPEG and collect bytes
compressed_frames = []
for frame in arr:
    img = Image.fromarray(frame)
    buffer = io.BytesIO()
    img.save(buffer, format='JPEG', quality=90)
    compressed_frames.append(buffer.getvalue())

# Send number of frames first, then each frame's length + data
sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
sock.connect(("receiver_ip", 12345))
sock.sendall(len(compressed_frames).to_bytes(4, byteorder='big'))
for frame_data in compressed_frames:
    sock.sendall(len(frame_data).to_bytes(4, byteorder='big'))
    sock.sendall(frame_data)
sock.close()

Receiver code:

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

sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
sock.bind(("0.0.0.0", 12345))
sock.listen(1)
conn, addr = sock.accept()

# Receive number of frames
num_frames = int.from_bytes(conn.recv(4), byteorder='big')

# Receive and decompress each frame
frames = []
for _ in range(num_frames):
    frame_len = int.from_bytes(conn.recv(4), byteorder='big')
    frame_data = b''
    while len(frame_data) < frame_len:
        chunk = conn.recv(min(4096, frame_len - len(frame_data)))
        if not chunk:
            break
        frame_data += chunk
    img = Image.open(io.BytesIO(frame_data))
    frames.append(np.array(img))

# Stack frames back into a 3D array
arr = np.stack(frames)
conn.close()
sock.close()

Key Implementation Tips #

• Avoid Sticky Packets: Always send the length of compressed data first, so the receiver knows exactly how many bytes to read. This prevents partial data or merged packets.
• Benchmark for Your Data: Test compression ratio and speed with your actual array—what works best depends on how much redundancy your data has. For example, zlib gives better compression than LZ4 but is slower.
• Chunked Transfer: For extremely large arrays, split the array into smaller chunks, compress each chunk, and send them sequentially. This reduces memory usage on both ends.

内容的提问来源于stack exchange，提问作者KaramJaber