131 lines
5.2 KiB
Python
Executable File
131 lines
5.2 KiB
Python
Executable File
'''
|
||
author : Hatcher W huimin
|
||
Email: hcwang@semi.ac.cn
|
||
use: 目标检测视频处理
|
||
input:需要处理的视频路径
|
||
output:avi后缀的视频处理结果
|
||
'''
|
||
import glob
|
||
import os
|
||
import time
|
||
|
||
import torch
|
||
from PIL import Image
|
||
from vizer.draw import draw_boxes
|
||
|
||
from ssd.config import cfg
|
||
from ssd.data.datasets import COCODataset, VOCDataset
|
||
import argparse
|
||
import numpy as np
|
||
|
||
from ssd.data.transforms import build_transforms
|
||
from ssd.modeling.detector import build_detection_model
|
||
from ssd.utils import mkdir
|
||
from ssd.utils.checkpoint import CheckPointer
|
||
from detection import main
|
||
from vediotoima import vedio_img
|
||
from imatovideo import images_to_video,images_to_video_big
|
||
from moviepy.editor import *
|
||
import os
|
||
os.environ['CUDA_VISIBLE_DEVICES'] = '2'
|
||
|
||
def parse_args():
|
||
|
||
parser = argparse.ArgumentParser(description='Detection Video to Video')
|
||
parser.add_argument(
|
||
"--config-file",
|
||
default="ssd/config/hat-512.yaml",
|
||
metavar="FILE",
|
||
help="path to config file",
|
||
type=str,
|
||
)
|
||
parser.add_argument("--ckpt", type=str, default="../data_MPKFNet/weigth/detection/detection.pth",
|
||
help="Trained weights.")
|
||
parser.add_argument("--score_threshold", type=float, default=0.3)
|
||
parser.add_argument("--video_name", default='000120210501080936.mp4', type=str, help='Video name.')
|
||
parser.add_argument("--images_dir", default='../data_MPKFNet/result_raw/yang/', type=str,
|
||
help='Specify a image dir to do prediction.')
|
||
parser.add_argument("--output_dir", default='../data_MPKFNet/result/yangd/', type=str,
|
||
help='Specify a image dir to save predicted images.')
|
||
parser.add_argument("--dataset_type", default="voc", type=str,
|
||
help='Specify dataset type. Currently support voc and coco.')
|
||
|
||
parser.add_argument(
|
||
"opts",
|
||
help="Modify config options using the command-line",
|
||
default=None,
|
||
nargs=argparse.REMAINDER,
|
||
)
|
||
args = parser.parse_args()
|
||
return args
|
||
|
||
|
||
# if __name__ == '__main__':
|
||
# args = parse_args()
|
||
# path = "/data/Hatcher/output1/"
|
||
# pathlists = os.listdir(path)
|
||
# for pathlist in pathlists:
|
||
# if os.path.exists('../data_MPKFNet/rawdata/' + pathlist + ".mp4"):
|
||
# print(pathlist," is skipping")
|
||
# continue
|
||
# print(pathlist," is starting")
|
||
# L = []
|
||
# for root, dirs, files in os.walk(os.path.join(path,pathlist)):
|
||
# files.sort()
|
||
# for file in files:
|
||
# if os.path.splitext(file)[1] in [".avi",".mp4"]:
|
||
# filepath = os.path.join(root, file)
|
||
# video = VideoFileClip(filepath)
|
||
# L.append(video)
|
||
# if len(L) == 0:
|
||
# continue
|
||
# final_clip = concatenate_videoclips(L)
|
||
# final_clip.to_videofile(os.path.join('../data_MPKFNet/rawdata/', "{}.mp4".format(pathlist)), fps=25, remove_temp=False)
|
||
# video_name_path = pathlist
|
||
# video_name_path = video_name_path.replace(".MP4", "").replace(".mp4", "").replace(".MOV", "").replace(".mov", "")
|
||
# RawVideoFilename = '../data_MPKFNet/rawdata/' + video_name_path + ".mp4"
|
||
# RawVideoExtrationSavedpath = '../data_MPKFNet/result_raw/' + video_name_path
|
||
# args.images_dir = '../data_MPKFNet/result_raw/' + video_name_path
|
||
# args.output_dir = "../data_MPKFNet/result/" + video_name_path + "Detection"
|
||
# CombineVideoRawImgPath = "../data_MPKFNet/result/" + video_name_path + "Detection"
|
||
# CombineVideoResultPath = '../data_MPKFNet/result_vedio11/' + "{}.avi".format(video_name_path)
|
||
#
|
||
# # vedio_img(RawVideoFilename, RawVideoExtrationSavedpath, 1)
|
||
# # main(args)
|
||
# # images_to_video(CombineVideoRawImgPath, CombineVideoResultPath)
|
||
# if os.path.exists(RawVideoExtrationSavedpath):
|
||
# print("skip video to image")
|
||
# pass
|
||
# else:
|
||
# vedio_img(RawVideoFilename, RawVideoExtrationSavedpath, 1)
|
||
# if os.path.exists(args.output_dir):
|
||
# print("skip deal")
|
||
# pass
|
||
# else:
|
||
# main(args)
|
||
# images_to_video(CombineVideoRawImgPath, CombineVideoResultPath)
|
||
|
||
|
||
#zhendui danyi shipin
|
||
if __name__ == '__main__':
|
||
args = parse_args()
|
||
video_name_path = args.video_name
|
||
video_name_path = video_name_path.replace(".MP4", "").replace(".mp4", "").replace(".MOV", "").replace(".mov", "")
|
||
RawVideoFilename = '../data_MPKFNet/rawdata/1/'+args.video_name
|
||
RawVideoExtrationSavedpath = '../data_MPKFNet/result_raw/' + video_name_path
|
||
args.images_dir = '../data_MPKFNet/result_raw/' + video_name_path
|
||
args.output_dir = "../data_MPKFNet/result/" + video_name_path + "Detection"
|
||
CombineVideoRawImgPath = "../data_MPKFNet/result/" + video_name_path + "Detection"
|
||
CombineVideoResultPath = '../data_MPKFNet/result_vedio/' + video_name_path + "Detection.avi"
|
||
if os.path.exists(RawVideoExtrationSavedpath):
|
||
print("skip video to image")
|
||
pass
|
||
else:
|
||
vedio_img(RawVideoFilename, RawVideoExtrationSavedpath, 1)
|
||
if os.path.exists(args.output_dir):
|
||
print("skip deal")
|
||
pass
|
||
else:
|
||
main(args)
|
||
images_to_video(CombineVideoRawImgPath, CombineVideoResultPath)
|