227 lines
7.8 KiB
Python
Executable File
227 lines
7.8 KiB
Python
Executable File
'''
|
||
author ： Hatcher W huimin
|
||
Email： hcwang@semi.ac.cn
|
||
use: 密度估计视频处理
|
||
input:视频路径
|
||
output：avi后缀的结果视频
|
||
'''
|
||
import torch
|
||
import datetime
|
||
import random
|
||
import cv2
|
||
import shutil
|
||
import glob
|
||
import json
|
||
import os
|
||
import numpy as np
|
||
from tqdm import tqdm
|
||
from configs.configs import *
|
||
from configs.configs import model_head
|
||
import matplotlib.pyplot as plt
|
||
from torch.utils.data import DataLoader
|
||
from torch.utils.data import Dataset
|
||
from dataset import People_Dataset
|
||
from network import CSRNet,MCNN,CANNet,PKSNet
|
||
import argparse
|
||
np.set_printoptions(threshold=np.inf)
|
||
|
||
# Path Definition
|
||
# RawVideoFilename = '../data_MPKFNet/rawdata/2021.2.20.mp4'
|
||
# RawVideoExtrationSavedpath = '../data_MPKFNet/result_raw/' + "2021.2.22"
|
||
# DensityEstimationResultPath = "../data_MPKFNet/result/2021.2.22midu"
|
||
# DensityEstimationRawImgPath = '../data_MPKFNet/result_raw/' + "2021.2.22"
|
||
# CombineVideoRawImgPath = "../data_MPKFNet/result/2021.2.22midu/"
|
||
# CombineVideoResultPath = '../data_MPKFNet/result_vedio/2021.2.22midu.avi'
|
||
ModelStateDict = '../data_MPKFNet/weigth/final.pth'
|
||
|
||
# Variable Definition
|
||
display = 5
|
||
get_name = []
|
||
|
||
'''
|
||
Cdoe Area
|
||
'''
|
||
# First Part--Video Frame Extrstion
|
||
def vedio_img(input,output,fps):
|
||
|
||
isExists = os.path.exists(output)
|
||
if not isExists:
|
||
os.makedirs(output)
|
||
print('path of %s is build' % (output))
|
||
else:
|
||
shutil.rmtree(output)
|
||
os.makedirs(output)
|
||
print('path of %s already exist and rebuild' % (output))
|
||
# 开始读视频
|
||
videoCapture = cv2.VideoCapture(input)
|
||
i = 0
|
||
j = 0
|
||
while True:
|
||
success, frame = videoCapture.read()
|
||
if not success:
|
||
print('video is all read')
|
||
break
|
||
i += 1
|
||
if (i % fps == 0):
|
||
# 保存图片
|
||
j += 1
|
||
savedname = output+"/" + "{}".format(j).zfill(7)+ '.jpg'
|
||
cv2.imwrite( savedname, frame)
|
||
print('image of %s is saved' % (savedname))
|
||
|
||
|
||
# Second Part--Density Etimation
|
||
class INf_Dataset(Dataset):
|
||
|
||
def __init__(self, DensityEstimationRawImgPath,img_list, train=True):
|
||
random.shuffle(img_list)
|
||
self.DensityEstimationRawImgPath=DensityEstimationRawImgPath
|
||
self.img_list = img_list
|
||
self.train = train
|
||
|
||
def __len__(self):
|
||
return len(self.img_list)
|
||
|
||
def __getitem__(self, index):
|
||
assert index <= len(self), 'index range error'
|
||
path_temp = self.DensityEstimationRawImgPath + "/" + self.img_list[index]
|
||
img = cv2.imread(path_temp) # H read image
|
||
img = cv2.resize(img, (1024, 768)) / 255.0
|
||
img = np.transpose(img, [2, 0, 1])
|
||
torch.from_numpy(img).type(torch.FloatTensor)
|
||
return img, self.img_list[index]
|
||
|
||
|
||
|
||
|
||
|
||
def inf(DensityEstimationResultPath,devices):
|
||
|
||
net = PKSNet().to(devices)
|
||
checkpoint = torch.load('../data_MPKFNet/weigth/final.pth')
|
||
net.load_state_dict(checkpoint['net'])
|
||
#net.load_state_dict(torch.load('../data_MPKFNet/weigth/sheep/bestmaesheep93,mse29262.pth'))
|
||
print("load model weight ok")
|
||
i = 0
|
||
for img, img_pathtemp in tqdm(data_loader):
|
||
|
||
img_pathtemp = img_pathtemp[0]
|
||
img = torch.tensor(img)
|
||
img = img.to(devices)
|
||
if model_head:
|
||
targets_pred, posture_pred, spice_pred = net(img)
|
||
else:
|
||
targets_pred = net(img)
|
||
t1 = datetime.datetime.now()
|
||
#num = targets_pred.data.sum()
|
||
targets_pred1 = targets_pred.cpu()
|
||
num = targets_pred1.data.sum().numpy()
|
||
#print(type(num))
|
||
targets_pred1 = targets_pred1.data.numpy()
|
||
|
||
t2 = datetime.datetime.now()
|
||
if i % display == 0:
|
||
num_display = num+15
|
||
data_dict = {"sum":str(num_display),"density":(str(targets_pred1[0, 0, :, :]))}
|
||
#print(data_dict)
|
||
json_outputpath = (DensityEstimationResultPath + "/" + str(img_pathtemp)).replace("/result/", "/result_json/").replace(".jpg",".json")
|
||
json_data = json.dumps(data_dict,ensure_ascii=False,indent=4)
|
||
with open (json_outputpath,'w',encoding='utf-8') as f:
|
||
f.write(str(json_data))
|
||
|
||
plt.figure(num_display)
|
||
plt.title("predict:{}".format(int(num_display)))
|
||
plt.imshow(targets_pred1[0, 0, :, :])
|
||
plt.savefig(DensityEstimationResultPath + "/" + str(img_pathtemp))
|
||
plt.close()
|
||
|
||
i += 1
|
||
|
||
|
||
# Three Part Combine Video
|
||
def images_to_video(path,CombineVideoResultPath):
|
||
|
||
img_array = []
|
||
img_list = [os.path.join(path,nm) for nm in os.listdir(path) if 1]
|
||
img_list.sort()
|
||
i=0
|
||
for filename in img_list:
|
||
print(i)
|
||
i+=1
|
||
img = cv2.imread(filename)
|
||
if img is None:
|
||
print(filename + " is error!")
|
||
continue
|
||
img_array.append(img)
|
||
# 图片的大小需要一致
|
||
#img_array, size = resize(img_array, 'largest')
|
||
fps = 25
|
||
out = cv2.VideoWriter(CombineVideoResultPath, cv2.VideoWriter_fourcc(*'DIVX'), fps, (640,480))
|
||
for i in range(len(img_array)):
|
||
out.write(img_array[i])
|
||
out.release()
|
||
|
||
if __name__ == '__main__':
|
||
|
||
# Build Ralted Saved Path
|
||
|
||
parser = argparse.ArgumentParser(description="DensityVideodemo.")
|
||
parser.add_argument("video_name", default='', type=str,
|
||
help='Video name.')
|
||
parser.add_argument("skip_getimg", default=False, type=str,
|
||
help='skip_getimg.')
|
||
parser.add_argument("GPU", default='6', type=str,
|
||
help='GPU.')
|
||
|
||
parser.add_argument(
|
||
"opts",
|
||
help="Modify config options using the command-line",
|
||
default=None,
|
||
nargs=argparse.REMAINDER,
|
||
)
|
||
args = parser.parse_args()
|
||
print(args)
|
||
# Gpu Selection
|
||
os.environ['CUDA_VISIBLE_DEVICES'] = args.GPU
|
||
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
|
||
video_name_path = args.video_name
|
||
video_name_path = video_name_path.replace(".MP4","").replace(".mp4","").replace(".MOV","").replace(".mov","")
|
||
RawVideoFilename = '../data_MPKFNet/rawdata/'+args.video_name
|
||
RawVideoExtrationSavedpath = '../data_MPKFNet/result_raw/' + video_name_path
|
||
DensityEstimationResultPath = "../data_MPKFNet/result/"+video_name_path+"desity"
|
||
DensityEstimationRawImgPath = '../data_MPKFNet/result_raw/' + video_name_path
|
||
CombineVideoRawImgPath = "../data_MPKFNet/result/"+video_name_path+"desity"
|
||
Json_ouputPath = "../data_MPKFNet/result_json/" + video_name_path + "desity"
|
||
CombineVideoResultPath = '../data_MPKFNet/result_vedio/'+video_name_path+"desity.avi"
|
||
print(args.skip_getimg)
|
||
print(args.GPU)
|
||
if args.skip_getimg == "False":
|
||
|
||
vedio_img(RawVideoFilename, RawVideoExtrationSavedpath, 1)
|
||
get_name = []
|
||
for _, _, filenames in os.walk(RawVideoExtrationSavedpath): # root是一个list是正在遍历的文件夹本身的地址 dirs文件夹中的目录 files是文件夹下的所有文件
|
||
for filename in filenames:
|
||
get_name.append(filename)
|
||
train = RawVideoExtrationSavedpath
|
||
trainfile = os.listdir(RawVideoExtrationSavedpath)
|
||
train_json = [file for file in trainfile]
|
||
train_json.sort()
|
||
with open(inf_json, 'w') as f:
|
||
json.dump(train_json, f)
|
||
with open(inf_json, 'r') as outfile:
|
||
inf_list = json.load(outfile)
|
||
inf_dataset = INf_Dataset(DensityEstimationRawImgPath,inf_list)
|
||
inf_dataset.img_list.sort()
|
||
data_loader = DataLoader(inf_dataset, 1)
|
||
|
||
|
||
if not os.path.exists(DensityEstimationResultPath) or not os.path.exists(DensityEstimationResultPath):
|
||
os.mkdir(DensityEstimationResultPath)
|
||
if not os.path.exists(Json_ouputPath) :
|
||
os.mkdir(Json_ouputPath)
|
||
#os.mkdir(RawVideoExtrationSavedpath)
|
||
# Video Extration Density Estimation Combine Image
|
||
|
||
inf(DensityEstimationResultPath,device)
|
||
images_to_video(CombineVideoRawImgPath,CombineVideoResultPath)
|