def handle(self, video_path: str) -> bool:
super(KerasHandler, self).handle(video_path)
video = VideoObject(video_path)
if self.preload:
video.load_frames()
# --- cutter ---
cutter = VideoCutter()
res = cutter.cut(video)
stable, unstable = res.get_range(threshold=0.98, offset=3)
# --- classify ---
cl = KerasClassifier()
if self.model_path:
logger.info("load existed pre-train model")
cl.load_model(self.model_path)
else:
data_home = res.pick_and_save(stable,
self.frame_count,
to_dir=self.result_path)
cl.train(data_home)
self.classifier_result = cl.classify(video, stable)
# --- draw ---
r = Reporter()
r.draw(self.classifier_result, report_path=self.result_report_path)
return True
def keras_train(
train_data_path: str,
model_path: str,
# options
epochs: int = 10,
target_size: str = "600x800",
overwrite: bool = False,
**kwargs,
):
from stagesepx.classifier.keras import KerasClassifier
# handle args
target_size: typing.Sequence[int] = [
int(each) for each in target_size.split("x")
]
cl = KerasClassifier(
# 轮数
epochs=epochs,
# 保证数据集的分辨率统一性
target_size=target_size,
**kwargs,
)
cl.train(train_data_path)
# file existed
while os.path.isfile(model_path):
logger.warning(f"file {model_path} already existed")
model_path = f"{uuid.uuid4()}.h5"
logger.debug(f"trying to save it to {model_path}")
cl.save_model(model_path, overwrite=overwrite)
def keras_train(
train_data_path: str,
model_path: str,
# options
epochs: int = 10,
target_size: str = "600x800",
overwrite: bool = False,
**kwargs,
):
from stagesepx.classifier.keras import KerasClassifier
assert not os.path.isfile(model_path), f"file {model_path} already existed"
# handle args
target_size: typing.Sequence[int] = [
int(each) for each in target_size.split("x")
]
cl = KerasClassifier(
# 轮数
epochs=epochs,
# 保证数据集的分辨率统一性
target_size=target_size,
**kwargs,
)
cl.train(train_data_path)
cl.save_model(model_path, overwrite=overwrite)
def train_model_Keras(_train_picture_path, _model_file_name):
cl = KerasClassifier(
# 轮数
epochs=10,
compress_rate=0.2,
# 保证数据集的分辨率统一性
# target_size=(600, 800),
)
cl.train(_train_picture_path)
cl.save_model(_model_file_name, overwrite=True)
return cl
def calculate_result(_cl, _SVM_or_Keras, _param, _from_movie_2_picture,
_model_file, _video_path_for_forecast):
if _SVM_or_Keras == '1\n':
print('使用SVM进行预测')
cl = SVMClassifier(
# 默认情况下使用 HoG 进行特征提取。你可以将其关闭从而直接对原始图片进行训练与测试：feature_type='raw'
feature_type="hog",
# 默认为0.2，即将图片缩放为0.2倍。主要为了提高计算效率,如果你担心影响分析效果，可以将其提高
compress_rate=_param[0],
)
# 加载待训练数据
cl.load_model(_model_file)
elif _SVM_or_Keras == '2\n':
# 分析视频_聚类
print('使用KerasClassifier进行预测')
cl = KerasClassifier(
compress_rate=_param[0],
# 在使用时需要保证数据集格式统一（与训练集）。因为 train_model.py 用了 600x800，所以这里设定成一样的
# target_size=(600, 800),
)
cl.load_model(_model_file)
# 开始预测
_forecast_result = []
# 获取forecast文件夹的mp4文件列表
forecast_video_list = public_fun.get_mp4file_name(_video_path_for_forecast)
for i in forecast_video_list:
# 分析视频_切割视频
stable = get_data.get_range('forecast', i, _param,
_from_movie_2_picture)
classify_result = cl.classify(i, stable, keep_data=True)
result_dict = classify_result.to_dict()
_forecast_result.append(
public_fun.write_result_to_local(i, _from_movie_2_picture,
result_dict, classify_result))
# _forecast_result = [['5.mp4', '3.161888888888889', '1',
# '4.155888888888889', '1', '8.040888888888889', '1']]
return _forecast_result
def test_keras():
# set epochs to 1 for quickly training (test only)
cl = KerasClassifier(epochs=1)
cl.train(CUTTER_RESULT_DIR)
cl.save_model("haha.h5")
# recreate
cl = KerasClassifier()
cl.load_model("haha.h5")
stable, _ = cutter_res.get_range()
classify_result = cl.classify(VIDEO_PATH, stable, keep_data=True)
assert classify_result.to_dict()
from stagesepx.cutter import VideoCutter
from stagesepx.classifier.keras import KerasClassifier
from stagesepx.reporter import Reporter
from stagesepx.video import VideoObject
video_path = "../videos/douyin1.mp4"
video = VideoObject(video_path)
video.load_frames()
# --- cutter ---
cutter = VideoCutter()
res = cutter.cut(video)
stable, unstable = res.get_range()
# --- classifier ---
cl = KerasClassifier(
# 在使用时需要保证数据集格式统一（与训练集）
# 因为 train_model.py 用了 600x800，所以这里设定成一样的
target_size=(600, 800), )
model_file = "./keras_model.h5"
cl.load_model(model_file)
classify_result = cl.classify(video, stable, keep_data=True)
result_dict = classify_result.to_dict()
# --- draw ---
r = Reporter()
r.draw(classify_result)
def create_model(self) -> Sequential:
# overwrite this method to design your own model structure!
# model = Sequential()
# ...
pass
def train(self, data_path: str, *_, **__):
# ...
pass
# or use the default one
# and then init it
# epochs=1 is just a example
cl = KerasClassifier(epochs=1)
# train model and save weights
cl.train(data_home)
cl.save_model("keras_model.h5")
# you would better reuse the trained model for less time cost
# keras model takes much more time than SVM
# cl.load_model("keras_model.h5")
classify_result = cl.classify(video, stable, keep_data=True)
result_dict = classify_result.to_dict()
# --- draw ---
r = Reporter()
r.draw(classify_result)
'''
* @Description : 开始根据图像数据训练模型及预测
* @Autor : Tommy
* @Date : 2021-02-21 18:51:11
* @LastEditors : Tommy
* @LastEditTime : 2021-02-26 14:46:49
'''
from stagesepx.classifier.keras import KerasClassifier
# 训练模型文件
cl = KerasClassifier(
# 设置这个数，需要保证小于总样本数，否则报错。
# nb_train_samples=30,
# 训练轮数
epochs=30)
cl.train('./stable_frame')
cl.save_model('./model.h5', overwrite=True)
def run(config: typing.Union[dict, str]):
"""
run with config
:param config: config file path, or a preload dict
:return:
"""
class _VideoUserConfig(BaseModel):
path: str
pre_load: bool = True
fps: int = None
class _CutterUserConfig(BaseModel):
threshold: float = None
frame_count: int = None
offset: int = None
limit: int = None
block: int = None
# common
compress_rate: float = None
target_size: typing.Tuple[int, int] = None
class _ClassifierType(Enum):
SVM = "svm"
KERAS = "keras"
class _ClassifierUserConfig(BaseModel):
boost_mode: bool = None
classifier_type: _ClassifierType = _ClassifierType.SVM
model: str = None
# common
compress_rate: float = None
target_size: typing.Tuple[int, int] = None
class _CalcOperatorType(Enum):
BETWEEN = "between"
DISPLAY = "display"
class _CalcOperator(BaseModel):
name: str
calc_type: _CalcOperatorType
args: dict = dict()
class _CalcUserConfig(BaseModel):
output: str = None
ignore_error: bool = None
operators: typing.List[_CalcOperator] = None
class _ExtraUserConfig(BaseModel):
save_train_set: str = None
class UserConfig(BaseModel):
output: str
video: _VideoUserConfig
cutter: _CutterUserConfig = _CutterUserConfig()
classifier: _ClassifierUserConfig = _ClassifierUserConfig()
calc: _CalcUserConfig = _CalcUserConfig()
extras: _ExtraUserConfig = _ExtraUserConfig()
if isinstance(config, str):
# path
config_path = pathlib.Path(config)
assert config_path.is_file(), f"no config file found in {config_path}"
# todo: support different types in the future
assert config_path.as_posix().endswith(
".json"), "config file should be json format"
with open(config_path, encoding=constants.CHARSET) as f:
config = json.load(f)
config = UserConfig(**config)
logger.info(f"config: {config}")
# main flow
video = VideoObject(
# fmt: off
path=config.video.path,
fps=config.video.fps,
)
if config.video.pre_load:
video.load_frames()
# cut
cutter = VideoCutter(
# fmt: off
compress_rate=config.cutter.compress_rate,
target_size=config.cutter.target_size,
)
res = cutter.cut(
# fmt: off
video=video,
block=config.cutter.block,
)
stable, unstable = res.get_range(
# fmt: off
threshold=config.cutter.threshold,
offset=config.cutter.offset,
)
with tempfile.TemporaryDirectory() as temp_dir:
# classify
if config.classifier.classifier_type is _ClassifierType.SVM:
cl = SVMClassifier(
# fmt: off
compress_rate=config.classifier.compress_rate,
target_size=config.classifier.target_size,
)
elif config.classifier.classifier_type is _ClassifierType.KERAS:
from stagesepx.classifier.keras import KerasClassifier
cl = KerasClassifier(
# fmt: off
compress_rate=config.classifier.compress_rate,
target_size=config.classifier.target_size,
)
# validation has been applied by pydantic
# so no `else`
if config.classifier.model:
# no need to retrain
model_path = pathlib.Path(config.classifier.model)
assert model_path.is_file(), f"file {model_path} not existed"
cl.load_model(model_path)
else:
# train a new model
train_set_dir = config.extras.save_train_set or temp_dir
os.makedirs(train_set_dir, exist_ok=True)
res.pick_and_save(
# fmt: off
stable,
frame_count=config.cutter.frame_count,
to_dir=train_set_dir,
)
cl.train(data_path=train_set_dir)
# start classifying
classify_result = cl.classify(
# fmt: off
video,
stable,
boost_mode=config.classifier.boost_mode,
)
# calc
def _calc_display() -> dict:
# jsonify
return json.loads(classify_result.dumps())
def _calc_between(*, from_stage: str = None, to_stage: str = None) -> dict:
assert classify_result.contain(
from_stage), f"no stage {from_stage} found in result"
assert classify_result.contain(
to_stage), f"no stage {to_stage} found in result"
from_frame = classify_result.last(from_stage)
to_frame = classify_result.first(to_stage)
cost = to_frame.timestamp - from_frame.timestamp
return {
"from": from_frame.frame_id,
"to": to_frame.frame_id,
"cost": cost,
}
_calc_func_dict = {
_CalcOperatorType.BETWEEN: _calc_between,
_CalcOperatorType.DISPLAY: _calc_display,
}
calc_output = config.calc.output
if calc_output:
output_path = pathlib.Path(calc_output)
assert not output_path.is_file(), f"file {output_path} already existed"
result = []
for each_calc in config.calc.operators:
func = _calc_func_dict[each_calc.calc_type]
try:
func_ret = func(**each_calc.args)
except Exception as e:
if not config.calc.ignore_error:
raise
logger.warning(e)
func_ret = traceback.format_exc()
calc_ret = {
"name": each_calc.name,
"type": each_calc.calc_type.value,
"result": func_ret,
}
result.append(calc_ret)
with open(output_path, "w", encoding=constants.CHARSET) as f:
json.dump(result, f)
# draw
r = Reporter()
r.draw(
# fmt: off
classify_result,
report_path=config.output,
)
'''
* @Description : 预测
* @Autor : Tommy
* @Date : 2021-02-21 19:07:26
* @LastEditors : Tommy
* @LastEditTime : 2021-03-15 12:15:48
'''
from stagesepx.classifier.keras import KerasClassifier
from stagesepx.cutter import VideoCutter
from stagesepx.video import VideoObject
from stagesepx.reporter import Reporter
# 使用Keras方法进行预测
cl = KerasClassifier()
cl.load_model('./Color_finally.h5')
# 将视频切分成帧
file_name = './video/TapColor_5.mp4'
video = VideoObject(file_name)
# 新建帧，计算视频总共多少帧，每帧多少ms
video.load_frames()
# 压缩视频
cutter = VideoCutter()
# 计算每一帧视频的每一个block的ssim和psnr
res = cutter.cut(video)
# 判断A帧到B帧之间是否稳定还是不稳定
stable, unstable = res.get_range()
from stagesepx.classifier.keras import KerasClassifier
data_home = "./dataset"
model_file = "./keras_model.h5"
cl = KerasClassifier(
# 轮数
epochs=10,
# 保证数据集的分辨率统一性
target_size=(600, 800),
)
cl.train(data_home)
cl.save_model(model_file, overwrite=True)
video.load_frames()
# 压缩视频
cutter = VideoCutter()
# 计算每一帧视频的每一个block的ssim和psnr。
res = cutter.cut(video, block=6)
# 计算出哪些区间是稳定的，哪些是不稳定的。判断A帧到B帧之间是稳定还是不稳定
stable, unstable = res.get_range(threshold=0.97, offset=2)
# 把分好类的稳定阶段的图片存本地
res.pick_and_save(stable,
100,
to_dir='./picture/train_stable_frame',
meaningful_name=True)
# 训练模型文件
cl = KerasClassifier(
# 训练轮数
epochs=10)
cl.train('./train_stable_frame')
cl.save_model('./model.h5', overwrite=True)
# 使用Keras方法进行预测
cl = KerasClassifier()
cl.load_model('./model.h5')
# 将视频切分成帧
file_name = './video_for_forecast.mp4'
# 预加载，大幅度提升分析速度
video = VideoObject(file_name, pre_load=True)
# 新建帧，计算视频总共有多少帧，每帧多少ms
video.load_frames()
# 压缩视频