Python---TCP Socket实现传送rk3588板子上的视频文件到电脑PC

服务端代码:传输视频

    
import time
import socket
import cv2
import numpy as np
import sys
from threading import Thread
from utils_2 import *
from rknnlite.api import RKNNLite
 
def get_cam(video):
    '''这是从电脑摄像机获取视频的函数'''
    capture = cv2.VideoCapture(video)
    frames_num=capture.get(7)
    print('视频总帧数：',frames_num)
    print('发送目标...')
    count = 0
    while capture.isOpened():
        ret, frames = capture.read()  # ret为返回值，frame为视频的每一帧
        yield frames
        #cv2.imshow('the video in local',frames)
        #c = cv2.waitKey(50)
        count += 1
    return frames
    
def post_cam(frame):
    '''这是将视频广播出去的函数'''
    frame=cv2.resize(frame,(360,270))  #控制画面大小
    img_encode = cv2.imencode('.jpg', frame)[1]
    data_encode = np.array(img_encode)
    data = data_encode.tobytes()
    print('数据处理成功')
    return data    

def segmentation(image_3c):
    # 预处理
    # image_2c = image_3c
    image_4c, image_3c = preprocess(image_3c, input_height=640, input_width=640)

    image_3C2 = image_3c[np.newaxis,:] #(1, 640, 640, 3)

    start = time.time()
    outputs = rknn_lite.inference(inputs=[image_3C2])
    stop = time.time()

    # 后处理
    results = postprocess(outputs, image_4c, image_3c, 0.4, 0.9, 2)
    results = results[0]              ## batch=1
    colorlist = [ (255,255,255), (0,0,0) ] 
    raw_img, seg_img, vis_img, right, left = vis_result(image_3c, results, colorlist, CLASSES, result_path)
    #cv2.imshow('seg_img', seg_img)
    #cv2.waitKey(2)   

    return raw_img, seg_img, vis_img, right, left        
            
            
if __name__ == "__main__":
    RKNN_MODEL = f'../model/green-640-640.rknn'
    CLASSES = ['road', 'lane_line']
    video = r'../3.mp4'
    result_path = f"../result"
    rknn_lite = RKNNLite()
    ret = rknn_lite.load_rknn(RKNN_MODEL)
    ret = rknn_lite.init_runtime()
    
    server_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    server_socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)  # 设置选项
    server_address = ('0.0.0.0',8003) #IP地址为服务端的IP地址，端口号任意
    server_socket.bind(server_address)
    server_socket.listen(1)
    print('Server is listening...')  
    client_socket, client_address = server_socket.accept()
    print('Connection from', client_address)
    
    # 1.读取图像
    frames= get_cam(video)  #调用本机的摄像机，只需一次，不用写入线程
    for frame in frames:
        # 2.进行语义分割处理
        start = time.time()
        rawimg2, segimg2, vis_img2, right2, left2 = segmentation(frame)
        end = time.time()
        print('车道线语义分割处理成功，耗时为 ', end - start, '秒')   
        vis_img = np.concatenate([rawimg2, segimg2, vis_img2], axis=1)  
        cv2.imshow('vis_img', vis_img)  
        cv2.waitKey(50)
        data = post_cam(segimg2)
        client_socket.sendall(data)
    print("发送成功")
    #client_socket.close()
    server_socket.close()

客户端代码:保存传输来的视频

import socket
import cv2
import sys
import time 
import datetime
import getpass
import numpy as np

                
if __name__ == '__main__':
    client_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    client_socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)  # 设置选项
    server_address = ('0.0.0.0',8003) #IP地址为服务端的IP地址，端口号任意
    client_socket.connect(server_address)
    imgs = []
    cv2.startWindowThread() 
    while True:
        data, addr = client_socket.recvfrom(400000000)
        nparr = np.frombuffer(data, np.uint8)
        img_decode = cv2.imdecode(nparr, cv2.IMREAD_COLOR)

        cv2.imshow('receive',img_decode)  #视频展示
        c = cv2.waitKey(2)
        if c == 27:  # 按了esc候可以退出
            break
        #一帧一帧保存图片成视频
        imgs.append(img_decode)
        v=cv2.VideoWriter(f'../result/cam.avi',\
            cv2.VideoWriter_fourcc(*'MJPG'),5,(360,270))
        for i in imgs:
            i = cv2.resize(i,(360,270))
            v.write(i)
    cv2.destroyAllWindows()
    #client_socket.close()
    

utils_2.py:

import cv2
import time
import numpy as np

def xywh2xyxy(x):
    y = np.copy(x)
    y[..., 0] = x[..., 0] - x[..., 2] / 2  # top left x
    y[..., 1] = x[..., 1] - x[..., 3] / 2  # top left y
    y[..., 2] = x[..., 0] + x[..., 2] / 2  # bottom right x
    y[..., 3] = x[..., 1] + x[..., 3] / 2  # bottom right y
    return y

def clip_boxes(boxes, shape):
    boxes[..., [0, 2]] = boxes[..., [0, 2]].clip(0, shape[1])  # x1, x2
    boxes[..., [1, 3]] = boxes[..., [1, 3]].clip(0, shape[0])  # y1, y2


def scale_boxes(img1_shape, boxes, img0_shape, ratio_pad=None):
    if ratio_pad is None:  # calculate from img0_shape
        gain = min(img1_shape[0] / img0_shape[0], img1_shape[1] / img0_shape[1])  # gain  = old / new
        pad = (img1_shape[1] - img0_shape[1] * gain) / 2, (img1_shape[0] - img0_shape[0] * gain) / 2  # wh padding
    else:
        gain = ratio_pad[0][0]
        pad = ratio_pad[1]

    boxes[..., [0, 2]] -= pad[0]  # x padding
    boxes[..., [1, 3]] -= pad[1]  # y padding
    boxes[..., :4] /= gain
    clip_boxes(boxes, img0_shape)
    return boxes

def crop_mask(masks, boxes):
    n, h, w = masks.shape
    x1, y1, x2, y2 = np.split(boxes[:, :, None], 4, axis=1)
    r = np.arange(w, dtype=np.float32)[None, None, :]  # rows shape(1,w,1)
    c = np.arange(h, dtype=np.float32)[None, :, None]  # cols shape(h,1,1)

    return masks * ((r >= x1) * (r < x2) * (c >= y1) * (c < y2))

def sigmoid(x): 
    return 1.0/(1+np.exp(-x))

def process_mask(protos, masks_in, bboxes, shape):

    c, mh, mw = protos.shape  # CHW
    ih, iw = shape
    masks = sigmoid(masks_in @ protos.reshape(c, -1)).reshape(-1, mh, mw)  # CHW 【lulu】

    downsampled_bboxes = bboxes.copy()
    downsampled_bboxes[:, 0] *= mw / iw
    downsampled_bboxes[:, 2] *= mw / iw
    downsampled_bboxes[:, 3] *= mh / ih
    downsampled_bboxes[:, 1] *= mh / ih

    masks = crop_mask(masks, downsampled_bboxes)  # CHW
    masks = np.transpose(masks, [1,2,0])
    # masks = cv2.resize(masks, (shape[1], shape[0]), interpolation=cv2.INTER_NEAREST)
    masks = cv2.resize(masks, (shape[1], shape[0]), interpolation=cv2.INTER_LINEAR)
    if masks.ndim == 3:
        masks = np.transpose(masks, [2,0,1])
    return np.where(masks>0.5,masks,0)

def nms(bboxes, scores, threshold=0.5):
    x1 = bboxes[:, 0]
    y1 = bboxes[:, 1]
    x2 = bboxes[:, 2]
    y2 = bboxes[:, 3]
    areas = (x2 - x1) * (y2 - y1)

    order = scores.argsort()[::-1]
    keep = []
    while order.size > 0:
        i = order[0]
        keep.append(i)

        if order.size == 1: break
        xx1 = np.maximum(x1[i], x1[order[1:]])
        yy1 = np.maximum(y1[i], y1[order[1:]])
        xx2 = np.minimum(x2[i], x2[order[1:]])
        yy2 = np.minimum(y2[i], y2[order[1:]])
        w = np.maximum(0.0, (xx2 - xx1))
        h = np.maximum(0.0, (yy2 - yy1))
        inter = w * h

        iou = inter / (areas[i] + areas[order[1:]] - inter)
        ids = np.where(iou <= threshold)[0]
        order = order[ids + 1]

    return keep


def non_max_suppression(
        prediction,
        conf_thres=0.25,
        iou_thres=0.45,
        classes=None,
        agnostic=False,
        multi_label=False,
        labels=(),
        max_det=300,
        nc=0,  # number of classes (optional)
):

    # Checks
    assert 0 <= conf_thres <= 1, f'Invalid Confidence threshold {conf_thres}, valid values are between 0.0 and 1.0'
    assert 0 <= iou_thres <= 1, f'Invalid IoU {iou_thres}, valid values are between 0.0 and 1.0'
    
    #【lulu】prediction.shape[1]：box + cls + num_masks
    bs = prediction.shape[0]              # batch size
    nc = nc or (prediction.shape[1] - 4)  # number of classes
    nm = prediction.shape[1] - nc - 4     # num_masks
    mi = 4 + nc                           # mask start index
    xc = np.max(prediction[:, 4:mi], axis=1) > conf_thres ## 【lulu】

    # Settings
    # min_wh = 2  # (pixels) minimum box width and height
    max_wh = 7680  # (pixels) maximum box width and height
    max_nms = 30000  # maximum number of boxes into torchvision.ops.nms()
    time_limit = 0.5 + 0.05 * bs  # seconds to quit after
    redundant = True  # require redundant detections
    multi_label &= nc > 1  # multiple labels per box (adds 0.5ms/img)
    merge = False  # use merge-NMS

    t = time.time()
    output = [np.zeros((0,6 + nm))] * bs ## 【lulu】

    for xi, x in enumerate(prediction):  # image_3c index, image_3c inference
        # Apply constraints
        # x[((x[:, 2:4] < min_wh) | (x[:, 2:4] > max_wh)).any(1), 4] = 0  # width-height

        x = np.transpose(x,[1,0])[xc[xi]] ## 【lulu】
        # If none remain process next image_3c
        if not x.shape[0]: continue

        # Detections matrix nx6 (xyxy, conf, cls)
        box, cls, mask = np.split(x, [4, 4+nc], axis=1) ## 【lulu】
        box = xywh2xyxy(box)  # center_x, center_y, width, height) to (x1, y1, x2, y2)

        j = np.argmax(cls, axis=1)  ## 【lulu】
        conf = cls[np.array(range(j.shape[0])), j].reshape(-1,1)
        x = np.concatenate([box, conf, j.reshape(-1,1), mask], axis=1)[conf.reshape(-1,)>conf_thres]

        # Check shape
        n = x.shape[0]  # number of boxes
        if not n: continue
        x = x[np.argsort(x[:, 4])[::-1][:max_nms]]  # sort by confidence and remove excess boxes 【lulu】
        # Batched NMS
        c = x[:, 5:6] * max_wh  # classes
        boxes, scores = x[:, :4] + c, x[:, 4]  # boxes (offset by class), scores
        i = nms(boxes, scores, iou_thres) ## 【lulu】
        i = i[:max_det]  # limit detections

        output[xi] = x[i]

        if (time.time() - t) > time_limit:
            # LOGGER.warning(f'WARNING ⚠️ NMS time limit {time_limit:.3f}s exceeded')
            break  # time limit exceeded
    return output

def make_anchors(feats_shape, strides, grid_cell_offset=0.5):
    """Generate anchors from features."""
    anchor_points, stride_tensor = [], []
    assert feats_shape is not None
    dtype_ = np.float
    for i, stride in enumerate(strides):
        _, _, h, w = feats_shape[i]
        sx = np.arange(w, dtype=dtype_) + grid_cell_offset  # shift x
        sy = np.arange(h, dtype=dtype_) + grid_cell_offset  # shift y

        sy, sx = np.meshgrid(sy, sx, indexing='ij') 
        anchor_points.append(np.stack((sx, sy), -1).reshape(-1, 2))
        stride_tensor.append(np.full((h * w, 1), stride, dtype=dtype_))
    return np.concatenate(anchor_points), np.concatenate(stride_tensor)

def dist2bbox(distance, anchor_points, xywh=True, dim=-1):
    """Transform distance(ltrb) to box(xywh or xyxy)."""
    lt, rb = np.split(distance, 2, dim)
    x1y1 = anchor_points - lt
    x2y2 = anchor_points + rb
    if xywh:
        c_xy = (x1y1 + x2y2) / 2
        wh = x2y2 - x1y1
        return np.concatenate((c_xy, wh), dim)  # xywh bbox
    return np.concatenate((x1y1, x2y2), dim)  # xyxy bbox

def letterbox(im, new_shape=(640, 640), color=(114, 114, 114), auto=True, scaleFill=False, scaleup=True, stride=32):
    # Resize and pad image while meeting stride-multiple constraints
    shape = im.shape[:2]  # current shape [height, width]
    if isinstance(new_shape, int):
        new_shape = (new_shape, new_shape)

    # Scale ratio (new / old)
    r = min(new_shape[0] / shape[0], new_shape[1] / shape[1])
    if not scaleup:  # only scale down, do not scale up (for better val mAP)
        r = min(r, 1.0)

    # Compute padding
    ratio = r, r  # width, height ratios
    new_unpad = int(round(shape[1] * r)), int(round(shape[0] * r))
    dw, dh = new_shape[1] - new_unpad[0], new_shape[0] - new_unpad[1]  # wh padding
    if auto:  # minimum rectangle
        dw, dh = np.mod(dw, stride), np.mod(dh, stride)  # wh padding
    elif scaleFill:  # stretch
        dw, dh = 0.0, 0.0
        new_unpad = (new_shape[1], new_shape[0])
        ratio = new_shape[1] / shape[1], new_shape[0] / shape[0]  # width, height ratios

    dw /= 2  # divide padding into 2 sides
    dh /= 2

    if shape[::-1] != new_unpad:  # resize
        im = cv2.resize(im, new_unpad, interpolation=cv2.INTER_LINEAR)
    top, bottom = int(round(dh - 0.1)), int(round(dh + 0.1))
    left, right = int(round(dw - 0.1)), int(round(dw + 0.1))
    im = cv2.copyMakeBorder(im, top, bottom, left, right, cv2.BORDER_CONSTANT, value=color)  # add border
    return im, ratio, (dw, dh)

def preprocess(image, input_height, input_width):
    image_3c = image

    # Convert the image_3c color space from BGR to RGB
    image_3c = cv2.cvtColor(image_3c, cv2.COLOR_BGR2RGB)

    # Resize the image_3c to match the input shape
    image_3c, ratio, dwdh = letterbox(image_3c, new_shape=[input_height, input_width], auto=False)

    # Normalize the image_3c data by dividing it by 255.0
    image_4c = np.array(image_3c) / 255.0

    # Transpose the image_3c to have the channel dimension as the first dimension
    image_4c = np.transpose(image_4c, (2, 0, 1))  # Channel first

    # Expand the dimensions of the image_3c data to match the expected input shape
    image_4c = np.expand_dims(image_4c, axis=0).astype(np.float32)

    image_4c = np.ascontiguousarray(image_4c)  # contiguous
    
    # Return the preprocessed image_3c data
    return image_4c, image_3c

def postprocess(preds, img, orig_img, OBJ_THRESH, NMS_THRESH, classes=None):
    p = non_max_suppression(preds[0],
                                OBJ_THRESH,
                                NMS_THRESH,
                                agnostic=False,
                                max_det=300,
                                nc=classes,
                                classes=None)        
    #print(p)                    
    results = []
    proto = preds[1] 
    for i, pred in enumerate(p):
        shape = orig_img.shape
        if not len(pred):
            results.append([[], [], []])  # save empty boxes
            continue
        masks = process_mask(proto[i], pred[:, 6:], pred[:, :4], img.shape[2:])  # HWC	
        # centroid = np.mean(np.argwhere(dummy_img),axis=0) 	
        pred[:, :4] = scale_boxes(img.shape[2:], pred[:, :4], shape).round()
        results.append([pred[:, :6], masks, shape[:2]])
    return results

def gen_color(class_num):
    color_list = []
    np.random.seed(1)
    while 1:
        a = list(map(int, np.random.choice(range(255),3)))
        if(np.sum(a)==0): continue
        color_list.append(a)
        if len(color_list)==class_num: break
    return color_list

def vis_result(image_3c, results, colorlist, CLASSES, result_path):
    boxes, masks, shape = results
    #if masks.ndim == 2:
    #    masks = np.expand_dims(masks, axis=0).astype(np.float32)
    image_3c = cv2.cvtColor(image_3c, cv2.COLOR_RGB2BGR)
    mask_img = np.zeros_like(image_3c)
    cls_list = []    
    center_list = []
    mask2 = np.zeros((640,640))
    mask2_img = np.zeros_like(image_3c)
    
    for box, mask in zip(boxes, masks):
        cls = int(box[-1])
        cls_list.append(cls)
        # 问题 一类时为（640，） 两类是为（640，640）
        # 解决 前面升维
        # print(mask.shape)
        dummy_img = np.zeros_like(image_3c)
        dummy_img[mask != 0] = colorlist[int(box[-1])]
        if cls == 0:
            mask_img[mask != 0] = (255, 255, 255)
        else :
            mask2 += mask
            mask_img[mask != 0] = (100, 100, 100)
        centroid = np.mean(np.argwhere(dummy_img), axis=0)
        if np.isnan(centroid).all() == False:
            centroid_x, centroid_y = int(centroid[1]), int(centroid[0])
            center_list.append([centroid_x, centroid_y])
    #image_3c, mask_img= image_3c[140:500, :], mask_img[140:500, :]
    mask2_img[mask2!=0] = (114, 114, 114)
    mask2_img = mask2_img[140:500,:]
    
    # 计算车道线距离
    min = 639
    max = 0
    for i, numbers in enumerate(mask2_img[200,:]):
        if numbers[0] != 0:
            if i < min:
                min = i
            elif i > max:
                max = i
    if max > 320 and min < 320:
        real_max = max-320
        real_min = 320-min
    elif max > 320 and min > 320:
        real_max = max-320
        real_min = 639
    elif max < 320 and min < 320:
        real_max = 639
        real_min = 320-min
    else:
        real_max = 639
        real_min = 639
        
    # 计算车道距离
    min2 = 639
    max2 = 0
    for i, numbers in enumerate(mask_img[200, :]):
        if numbers[0] != 0:
            if i < min2:
                min2 = i
            elif i>max2:
                max2 = i
	    
    if real_max == 999 and max2 > 320:
        real_max = max2-320
    if real_min == 999 and min2 < 320:
        real_min = 320-min2
    
    # 画点
    vis_img = cv2.addWeighted(image_3c, 0.5, mask_img, 0.5, 0)
    
    for j in range(len(center_list)):
        cv2.circle(vis_img, (center_list[j][0], center_list[j][1]), radius=5, color=(0, 0, 255), thickness=-1)
        if real_max<999:
            vis_img[200:205, real_max-5+320:real_max+320] = (0,0,255)
        if real_min<999:
            mask_img[200:205, 320-real_min:320-real_min+5] = (0,0,255)
        mask_img[200:205, 317:322] = (0,0,255)  
    vis_img2 = np.concatenate([image_3c, mask_img, vis_img],axis=1)
    
        
    cv2.imwrite(f"./{result_path}/origin_image.jpg", image_3c)
    cv2.imwrite(f"./{result_path}/mask_image.jpg", mask_img)
    cv2.imwrite(f"./{result_path}/visual_image.jpg", vis_img2)
    
    return image_3c, mask_img, vis_img, real_max, real_min