''' 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)