298 lines
11 KiB
Python
298 lines
11 KiB
Python
import copy
|
|
import subprocess as sp
|
|
from enum import Enum, unique
|
|
from PIL import Image
|
|
import time
|
|
import cv2
|
|
|
|
import sys
|
|
sys.path.extend(['..','../AIlib' ])
|
|
|
|
from AI import AI_process, AI_process_forest, get_postProcess_para
|
|
import cv2,os,time
|
|
from segutils.segmodel import SegModel
|
|
from models.experimental import attempt_load
|
|
from utils.torch_utils import select_device
|
|
from utilsK.queRiver import get_labelnames,get_label_arrays
|
|
import numpy as np
|
|
import torch
|
|
from utilsK.masterUtils import get_needed_objectsIndex
|
|
|
|
|
|
|
|
|
|
|
|
# 异常枚举
|
|
@unique
|
|
class ModelType(Enum):
|
|
|
|
WATER_SURFACE_MODEL = ("1", "001", "水面模型")
|
|
|
|
FOREST_FARM_MODEL = ("2", "002", "森林模型")
|
|
|
|
TRAFFIC_FARM_MODEL = ("3", "003", "交通模型")
|
|
|
|
def checkCode(code):
|
|
for model in ModelType:
|
|
if model.value[1] == code:
|
|
return True
|
|
return False
|
|
class ModelConfig():
|
|
def __init__(self):
|
|
postFile = '../AIlib/conf/para.json'
|
|
self.conf_thres, self.iou_thres, self.classes, self.rainbows = get_postProcess_para(postFile)
|
|
|
|
|
|
class SZModelConfig(ModelConfig):
|
|
|
|
def __init__(self):
|
|
super(SZModelConfig, self).__init__()
|
|
labelnames = "../AIlib/weights/yolov5/class8/labelnames.json" ##对应类别表
|
|
self.names = get_labelnames(labelnames)
|
|
self.label_arraylist = get_label_arrays(self.names, self.rainbows, outfontsize=40,
|
|
fontpath="../AIlib/conf/platech.ttf")
|
|
|
|
|
|
class LCModelConfig(ModelConfig):
|
|
def __init__(self):
|
|
super(LCModelConfig, self).__init__()
|
|
labelnames = "../AIlib/weights/forest/labelnames.json"
|
|
self.names = get_labelnames(labelnames)
|
|
self.label_arraylist = get_label_arrays(self.names, self.rainbows, outfontsize=40, fontpath="../AIlib/conf/platech.ttf")
|
|
|
|
|
|
class RFModelConfig(ModelConfig):
|
|
def __init__(self):
|
|
super(RFModelConfig, self).__init__()
|
|
labelnames = "../AIlib/weights/road/labelnames.json"
|
|
self.names = get_labelnames(labelnames)
|
|
imageW = 1536
|
|
outfontsize=int(imageW/1920*40)
|
|
self.label_arraylist = get_label_arrays(self.names, self.rainbows, outfontsize=outfontsize, fontpath="../AIlib/conf/platech.ttf")
|
|
|
|
class Model():
|
|
def __init__(self, device, allowedList=None):
|
|
##预先设置的参数
|
|
self.device_ = device ##选定模型，可选 cpu,'0','1'
|
|
self.allowedList = allowedList
|
|
|
|
|
|
# 水面模型
|
|
class SZModel(Model):
|
|
def __init__(self, device, allowedList=None):
|
|
super().__init__(device, allowedList)
|
|
|
|
self.device = select_device(self.device_)
|
|
self.half = self.device.type != 'cpu'
|
|
self.model = attempt_load("../AIlib/weights/yolov5/class8/bestcao.pt", map_location=self.device)
|
|
if self.half:
|
|
self.model.half()
|
|
self.segmodel = SegModel(nclass=2, weights='../AIlib/weights/STDC/model_maxmIOU75_1720_0.946_360640.pth',
|
|
device=self.device)
|
|
|
|
# names, label_arraylist, rainbows, conf_thres, iou_thres
|
|
def process(self, frame, config):
|
|
return AI_process([frame], self.model, self.segmodel, config[0], config[1],
|
|
config[2], self.half, self.device, config[3], config[4],
|
|
self.allowedList)
|
|
|
|
|
|
# 森林模型
|
|
class LCModel(Model):
|
|
def __init__(self, device, allowedList=None):
|
|
super().__init__(device, allowedList)
|
|
self.device = select_device(self.device_)
|
|
self.half = self.device.type != 'cpu' # half precision only supported on CUDA
|
|
self.model = attempt_load("../AIlib/weights/forest/best.pt", map_location=self.device) # load FP32 model
|
|
if self.half:
|
|
self.model.half()
|
|
self.segmodel = None
|
|
|
|
# names, label_arraylist, rainbows, conf_thres, iou_thres
|
|
def process(self, frame, config):
|
|
return AI_process_forest([frame], self.model, self.segmodel, config[0], config[1], config[2],
|
|
self.half, self.device, config[3], config[4], self.allowedList)
|
|
|
|
# 交通模型
|
|
class RFModel(Model):
|
|
def __init__(self, device, allowedList=None):
|
|
super().__init__(device, allowedList)
|
|
self.device = select_device(self.device_)
|
|
self.half = self.device.type != 'cpu' # half precision only supported on CUDA
|
|
self.model = attempt_load("../AIlib/weights/road/best.pt", map_location=self.device) # load FP32 model
|
|
if self.half:
|
|
self.model.half()
|
|
self.segmodel = None
|
|
|
|
# names, label_arraylist, rainbows, conf_thres, iou_thres
|
|
def process(self, frame, config):
|
|
return AI_process_forest([frame], self.model, self.segmodel, config[0], config[1], config[2],
|
|
self.half, self.device, config[3], config[4], self.allowedList)
|
|
|
|
def get_model(args):
|
|
for model in args[2]:
|
|
try:
|
|
code = '001'
|
|
needed_objectsIndex = [int(category.get("id")) for category in model.get("categories")]
|
|
if code == ModelType.WATER_SURFACE_MODEL.value[1]:
|
|
return SZModel(args[1], needed_objectsIndex), code, args[0].get("sz")
|
|
elif code == ModelType.FOREST_FARM_MODEL.value[1]:
|
|
return LCModel(args[1], needed_objectsIndex), code, args[0].get("lc")
|
|
elif code == ModelType.TRAFFIC_FARM_MODEL.value[1]:
|
|
return RFModel(args[1], needed_objectsIndex), code, args[0].get("rf")
|
|
else:
|
|
raise Exception("11111")
|
|
except Exception as e:
|
|
raise Exception("22222")
|
|
class PictureWaterMark():
|
|
|
|
def common_water(self, image, logo):
|
|
width, height = image.shape[1], image.shape[0]
|
|
mark_width, mark_height = logo.shape[1], logo.shape[0]
|
|
rate = int(width * 0.2) / mark_width
|
|
logo_new = cv2.resize(logo, None, fx=rate, fy=rate, interpolation=cv2.INTER_NEAREST)
|
|
position = (int(width * 0.95 - logo_new.shape[1]), int(height * 0.95 - logo_new.shape[0]))
|
|
b = Image.new('RGBA', (width, height), (0, 0, 0, 0)) # 创建新图像：透明'
|
|
a = Image.fromarray(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
|
|
watermark = Image.fromarray(cv2.cvtColor(logo_new, cv2.COLOR_BGRA2RGBA))
|
|
# 图片旋转
|
|
# watermark = watermark.rotate(45)
|
|
b.paste(a, (0, 0))
|
|
b.paste(watermark, position, mask=watermark)
|
|
return cv2.cvtColor(np.asarray(b), cv2.COLOR_BGR2RGB)
|
|
|
|
def common_water_1(self, image, logo, alpha=1):
|
|
h, w = image.shape[0], image.shape[1]
|
|
if w >= h:
|
|
rate = int(w * 0.1) / logo.shape[1]
|
|
else:
|
|
rate = int(h * 0.1) / logo.shape[0]
|
|
mask = cv2.resize(logo, None, fx=rate, fy=rate, interpolation=cv2.INTER_NEAREST)
|
|
mask_h, mask_w = mask.shape[0], mask.shape[1]
|
|
mask_channels = cv2.split(mask)
|
|
dst_channels = cv2.split(image)
|
|
# b, g, r, a = cv2.split(mask)
|
|
# 计算mask在图片的坐标
|
|
ul_points = (int(h * 0.95) - mask_h, int(w - h * 0.05 - mask_w))
|
|
dr_points = (int(h * 0.95), int(w - h * 0.05))
|
|
for i in range(3):
|
|
dst_channels[i][ul_points[0]: dr_points[0], ul_points[1]: dr_points[1]] = dst_channels[i][
|
|
ul_points[0]: dr_points[0],
|
|
ul_points[1]: dr_points[1]] * (
|
|
255.0 - mask_channels[3] * alpha) / 255
|
|
dst_channels[i][ul_points[0]: dr_points[0], ul_points[1]: dr_points[1]] += np.array(
|
|
mask_channels[i] * (mask_channels[3] * alpha / 255), dtype=np.uint8)
|
|
dst_img = cv2.merge(dst_channels)
|
|
return dst_img
|
|
def video_merge(frame1, frame2, width, height):
|
|
frameLeft = cv2.resize(frame1, (width, height), interpolation=cv2.INTER_LINEAR)
|
|
frameRight = cv2.resize(frame2, (width, height), interpolation=cv2.INTER_LINEAR)
|
|
frame_merge = np.hstack((frameLeft, frameRight))
|
|
# frame_merge = np.hstack((frame1, frame2))
|
|
return frame_merge
|
|
cap = cv2.VideoCapture("/home/DATA/chenyukun/4.mp4")
|
|
|
|
# Get video information
|
|
fps = int(cap.get(cv2.CAP_PROP_FPS))
|
|
print(fps)
|
|
width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
|
|
print(width)
|
|
height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
|
|
print(height)
|
|
# command = ['ffmpeg',
|
|
# '-y', # 不经过确认，输出时直接覆盖同名文件。
|
|
# '-f', 'rawvideo',
|
|
# '-vcodec', 'rawvideo',
|
|
# '-pix_fmt', 'bgr24',
|
|
# # '-s', "{}x{}".format(self.width * 2, self.height),
|
|
# '-s', "{}x{}".format(width, height),
|
|
# '-r', str(15),
|
|
# '-i', '-', # 指定输入文件
|
|
# '-g', '15',
|
|
# '-sc_threshold', '0', # 使得GOP的插入更加均匀
|
|
# '-b:v', '3000k', # 指定码率
|
|
# '-tune', 'zerolatency', # 加速编码速度
|
|
# '-c:v', 'libx264', # 指定视频编码器
|
|
# '-pix_fmt', 'yuv420p',
|
|
# "-an",
|
|
# '-preset', 'ultrafast', # 指定输出的视频质量，会影响文件的生成速度，有以下几个可用的值 ultrafast,
|
|
# # superfast, veryfast, faster, fast, medium, slow, slower, veryslow。
|
|
# '-f', 'flv',
|
|
# "rtmp://live.push.t-aaron.com/live/THSAk"]
|
|
#
|
|
# # 管道配置
|
|
# p = sp.Popen(command, stdin=sp.PIPE, shell=False)
|
|
|
|
sz = SZModelConfig()
|
|
lc = LCModelConfig()
|
|
rf = RFModelConfig()
|
|
config = {
|
|
"sz": (sz.names, sz.label_arraylist, sz.rainbows, sz.conf_thres, sz.iou_thres),
|
|
"lc": (lc.names, lc.label_arraylist, lc.rainbows, lc.conf_thres, lc.iou_thres),
|
|
"rf": (rf.names, rf.label_arraylist, rf.rainbows, rf.conf_thres, rf.iou_thres),
|
|
}
|
|
model = {
|
|
"models": [
|
|
{
|
|
"code": "001",
|
|
"categories": [
|
|
{
|
|
"id": "0",
|
|
"config": {}
|
|
},
|
|
{
|
|
"id": "1",
|
|
"config": {}
|
|
},
|
|
{
|
|
"id": "2",
|
|
"config": {}
|
|
},
|
|
{
|
|
"id": "3",
|
|
"config": {}
|
|
},
|
|
{
|
|
"id": "4",
|
|
"config": {}
|
|
},
|
|
{
|
|
"id": "5",
|
|
"config": {}
|
|
},
|
|
{
|
|
"id": "6",
|
|
"config": {}
|
|
},
|
|
{
|
|
"id": "7",
|
|
"config": {}
|
|
}
|
|
]
|
|
}]
|
|
}
|
|
|
|
mod, model_type_code, modelConfig = get_model((config, str(0), model.get("models")))
|
|
pic = PictureWaterMark()
|
|
logo = cv2.imread("./image/logo.png", -1)
|
|
ai_video_file = cv2.VideoWriter("/home/DATA/chenyukun/aa/2.mp4", cv2.VideoWriter_fourcc(*'mp4v'), fps, (width*2, height))
|
|
while(cap.isOpened()):
|
|
start =time.time()
|
|
ret, frame = cap.read()
|
|
# cap.grab()
|
|
if not ret:
|
|
print("Opening camera is failed")
|
|
break
|
|
p_result, timeOut = mod.process(copy.deepcopy(frame), modelConfig)
|
|
frame = pic.common_water_1(frame, logo)
|
|
p_result[1] = pic.common_water_1(p_result[1], logo)
|
|
frame_merge = video_merge(frame, p_result[1], width, height)
|
|
ai_video_file.write(frame_merge)
|
|
print(time.time()-start)
|
|
ai_video_file.release()
|
|
cap.release()
|
|
|
|
|
|
|
|
|