简单来说,就是pc端通过文生图、图生图等形式生成图片,通过socket传递给Aidlux端,Aidlux进行测评并返回结果。
下面简要介绍下Aidlux:Aidlux是一套基于安卓平台边缘计算系统,类似于在安卓平台中嵌入ubuntu系统,可以在安卓手机中下载,也可以在Aidlux发布的边缘计算盒子中运行,我是在Aidlux s855边缘计算盒子中使用,用安卓手机也可以运行demo,只不过边缘计算盒子考虑工业性,散热、接口更丰富。
本文采用的是本地gpu运行(win系统下运行,ubuntu系统下也可)
Aidlux端代码如下,在Aidlux端使用:python Socket_Aidlux.py 调用(注意,先运行Aidlux端代码,再运行pc端代码)
import socket
import time
import cv2
import numpy
import copy
# aidlux相关
from cvs import *
import aidlite_gpu
from utils import detect_postprocess, preprocess_img, draw_detect_res
#, extract_detect_res
import time
import cv2
def ReceiveVideo():
# IP地址'0.0.0.0'为等待客户端连接
address = ('192.168.31.5', 9000)
# 建立socket对象,参数意义见https://blog.csdn.net/rebelqsp/article/details/22109925
# socket.AF_INET:服务器之间网络通信
# socket.SOCK_STREAM:流式socket , for TCP
s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
# 将套接字绑定到地址, 在AF_INET下,以元组(host,port)的形式表示地址.
s.bind(address)
# 开始监听TCP传入连接。参数指定在拒绝连接之前,操作系统可以挂起的最大连接数量。该值至少为1,大部分应用程序设为5就可以了。
s.listen(5)
def recvall(sock, count):
buf = b'' # buf是一个byte类型
while count:
# 接受TCP套接字的数据。数据以字符串形式返回,count指定要接收的最大数据量.
newbuf = sock.recv(count)
if not newbuf: return None
buf += newbuf
count -= len(newbuf)
return buf
# 接受TCP连接并返回(conn,address),其中conn是 新的套接字对象,可以用来接收和发送数据。addr是连接客户端的地址。
# 没有连接则等待有连接
conn, addr = s.accept()
print('connect from PC:' + str(addr))
if 1:
start = time.time() # 用于计算帧率信息
length = recvall(conn, 16) # 获得图片文件的长度,16代表获取长度
stringData = recvall(conn, int(length)) # 根据获得的文件长度,获取图片文件
data = numpy.frombuffer(stringData, numpy.uint8) # 将获取到的字符流数据转换成1维数组
decimg = cv2.imdecode(data, cv2.IMREAD_COLOR) # 将数组解码成图像
cv2.imwrite("car.jpg",decimg)
print("save image ")
# cv2.imshow('SERVER', decimg) # 显示图像
# cv2.waitKey(2000)
#
# # 进行下一步处理
# # 。
# # 。
# # 。
#
# AidLite初始化:调用AidLite进行AI模型的加载与推理,需导入aidlite
aidlite = aidlite_gpu.aidlite()
# Aidlite模型路径
model_path = '/home//models/yolov5n-fp16.tflite'
# 定义输入输出shape
in_shape = [1 * 640 * 640 * 3 * 4]
out_shape = [1 * 25200 * 85 * 4]
# 加载Aidlite检测模型:支持tflite, tnn, mnn, ms, nb格式的模型加载
aidlite.ANNModel(model_path, in_shape, out_shape, 4, 0)
# 读取图片进行推理
# 设置测试集路径
source = "/home/AIGC/images/AIGC"
images_list = os.listdir(source)
print(images_list)
frame_id = 0
# 读取数据集
for image_name in images_list:
frame_id += 1
print("frame_id:", frame_id)
image_path = os.path.join(source, image_name)
frame = cvs.imread(image_path)
# 预处理
img = preprocess_img(frame, target_shape=(640, 640), div_num=255, means=None, stds=None)
# 数据转换:因为setTensor_Fp32()需要的是float32类型的数据,所以送入的input的数据需为float32,大多数的开发者都会忘记将图像的数据类型转换为float32
aidlite.setInput_Float32(img, 640, 640)
# 模型推理API
aidlite.invoke()
# 读取返回的结果
pred = aidlite.getOutput_Float32(0)
# 数据维度转换
pred = pred.reshape(1, 25200, 85)[0]
# 模型推理后处理
pred = detect_postprocess(pred, frame.shape, [640, 640, 3], conf_thres=0.25, iou_thres=0.45)
# 绘制推理结果
res_img = draw_detect_res(frame, pred)
cvs.imshow(res_img)
# 测试结果展示停顿
time.sleep(5)
# # 将帧率信息回传,主要目的是测试可以双向通信
end = time.time()
seconds = end - start
fps = 1 / seconds
##返回已处理图像到客户端
conn.send(bytes(str(int(fps)), encoding='utf-8'))
# image = copy.deepcopy(decimg)
# encode_param = [int(cv2.IMWRITE_JPEG_QUALITY), 100]
# result, imgencode = cv2.imencode('.jpg', image, encode_param)
# # 建立矩阵
# data = numpy.array(imgencode)
# # 将numpy矩阵转换成字符形式,以便在网络中传输
# img_Data = data.tostring()
# # 先发送要发送的数据的长度
# # ljust() 方法返回一个原字符串左对齐,并使用空格填充至指定长度的新字符串
# conn.send(str.encode(str(len(img_Data)).ljust(16)))
# # # print(img_Data)
# # # 发送数据
# conn.send(img_Data)
# if cv2.waitKey(10) & 0xff == 27:
# break
s.close()
# cv2.destroyAllWindows()
if __name__ == '__main__':
ReceiveVideo()
pc端代码:
import socket
import cv2
import numpy as np
import time
import sys
### 本代码主要是客户端代码,aidlux上的Socket_fuwuduan.py是匹配的服务端代码,当服务端代码启动时,由本代码读取一张图片,推送过去
import torch
# from torch import autocast
from torch.cuda.amp import autocast
from diffusers import StableDiffusionPipeline, DPMSolverMultistepScheduler,DDIMScheduler
def recvall(sock, count):
buf = b'' # buf是一个byte类型
while count:
newbuf = sock.recv(count)
if not newbuf: return None
buf += newbuf
count -= len(newbuf)
return buf
def SendAIGC():
# 建立sock连接
# address要连接的aidlux服务器IP地址和端口号
address = ('192.168.31.5', 9529)
try:
# 建立socket对象
# socket.AF_INET:服务器之间网络通信
# socket.SOCK_STREAM:流式socket , for TCP
sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
# 开启连接
sock.connect(address)
except socket.error as msg:
print(msg)
sys.exit(1)
###########传送AIGC图片#################
# ## 如果本地没有GPU
# if 1:
# frame = cv2.imread("car.png")
# # # 压缩参数,后面cv2.imencode将会用到,对于jpeg来说,15代表图像质量,越高代表图像质量越好为 0-100,默认95
# encode_param = [int(cv2.IMWRITE_JPEG_QUALITY), 95]
# # cv2.imencode将图片格式转换(编码)成流数据,赋值到内存缓存中;主要用于图像数据格式的压缩,方便网络传输
# # '.jpg'表示将图片按照jpg格式编码。
# result, imgencode = cv2.imencode('.jpg', frame, encode_param)
# # 建立矩阵
# data = np.array(imgencode)
# # 将numpy矩阵转换成字符形式,以便在网络中传输
# stringData = data.tostring()
#
# # 先发送要发送的数据的长度
# # ljust() 方法返回一个原字符串左对齐,并使用空格填充至指定长度的新字符串
# sock.send(str.encode(str(len(stringData)).ljust(16)))
# # 发送数据
# sock.send(stringData)
### 如果本地有GPU
# if 0:
### 本地生成AIGC图片 ###
## 添加AIGC代码 ##
#####################
model_id = "E:/pycharmwenjian/Aidlux课程/AIGC训练营_Lesson2_code/models/stable-diffusion-v1-5"
pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16)
# In[13]:
# pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
pipe.scheduler = DDIMScheduler.from_config(pipe.scheduler.config)
pipe = pipe.to("cuda")
# In[18]:
prompt = "a car"
negative_prompt = "(EasyNegative), bad_prompt_version2, (watermark), (signature), (sketch by bad-artist), (signature), (worst quality), (low quality), ((badhandsv5-neg)), ((badhandv4)), (bad anatomy), deformed hands, NSFW, nude, EasyNegative, (worst quality:1.4), (low quality:1.4), (normal quality:1.4),lowres,crowd"
# with autocast("cuda"):
image = pipe(prompt,
negative_prompt=negative_prompt,
num_inference_steps=25,
width=512,
height=768,
guidance_scale=7.5).images[0]
# In[19]:
image.save("./car.png")
frame = cv2.imread("./car.png")
# # 压缩参数,后面cv2.imencode将会用到,对于jpeg来说,15代表图像质量,越高代表图像质量越好为 0-100,默认95
encode_param = [int(cv2.IMWRITE_JPEG_QUALITY), 95]
# cv2.imencode将图片格式转换(编码)成流数据,赋值到内存缓存中;主要用于图像数据格式的压缩,方便网络传输
# '.jpg'表示将图片按照jpg格式编码。
result, imgencode = cv2.imencode('.jpg', frame, encode_param)
# 建立矩阵
data = np.array(imgencode)
# 将numpy矩阵转换成字符形式,以便在网络中传输
stringData = data.tostring()
# 先发送要发送的数据的长度
# ljust() 方法返回一个原字符串左对齐,并使用空格填充至指定长度的新字符串
sock.send(str.encode(str(len(stringData)).ljust(16)))
# 发送数据
sock.send(stringData)
# 读取服务器返回值
receive = sock.recv(16)
if len(receive):
print("图片发送成功")
print(str(receive, encoding='utf-8')) ### 之前接受的帧率数据,现在换成image流数据
sock.close()
if __name__ == '__main__':
SendAIGC()
运行结果:
推理速度、准确度,检测速度,开辟内存存储图片减少硬盘存、读,批量生成图片并检测都是本代码可拓展方向