first commit
This commit is contained in:
Semi212 committed
commit
54ccfb11fe
218 files changed
+19525
No files matched your search
Executable
+226
@@ -0,0 +1,226 @@
|
||||
'''
|
||||
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='0', 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 not args.skip_getimg:
|
||||
|
||||
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)
|
||||
Reference in new issue
Block a user