Initial clean project import
This commit is contained in:
cat-shark
@@ -0,0 +1,97 @@
|
||||
# Temporary directories and files
|
||||
*.ckpt
|
||||
*.onnx
|
||||
|
||||
# Byte-compiled / optimized / DLL files
|
||||
__pycache__/
|
||||
*.py[cod]
|
||||
*$py.class
|
||||
|
||||
# C extensions
|
||||
*.so
|
||||
|
||||
# Distribution / packaging
|
||||
.Python
|
||||
env/
|
||||
build/
|
||||
develop-eggs/
|
||||
dist/
|
||||
downloads/
|
||||
eggs/
|
||||
.eggs/
|
||||
lib/
|
||||
lib64/
|
||||
parts/
|
||||
sdist/
|
||||
var/
|
||||
*.egg-info/
|
||||
.installed.cfg
|
||||
*.egg
|
||||
|
||||
# PyInstaller
|
||||
# Usually these files are written by a python script from a template
|
||||
# before PyInstaller builds the exe, so as to inject date/other infos into it.
|
||||
*.manifest
|
||||
*.spec
|
||||
|
||||
# Installer logs
|
||||
pip-log.txt
|
||||
pip-delete-this-directory.txt
|
||||
|
||||
# Unit test / coverage reports
|
||||
htmlcov/
|
||||
.tox/
|
||||
.coverage
|
||||
.coverage.*
|
||||
.cache
|
||||
nosetests.xml
|
||||
coverage.xml
|
||||
*,cover
|
||||
.hypothesis/
|
||||
|
||||
# Translations
|
||||
*.mo
|
||||
*.pot
|
||||
|
||||
# Django stuff:
|
||||
*.log
|
||||
local_settings.py
|
||||
|
||||
# Flask stuff:
|
||||
instance/
|
||||
.webassets-cache
|
||||
|
||||
# Scrapy stuff:
|
||||
.scrapy
|
||||
|
||||
# Sphinx documentation
|
||||
docs/_build/
|
||||
|
||||
# PyBuilder
|
||||
target/
|
||||
|
||||
# IPython Notebook
|
||||
.ipynb_checkpoints
|
||||
|
||||
# pyenv
|
||||
.python-version
|
||||
|
||||
# celery beat schedule file
|
||||
celerybeat-schedule
|
||||
|
||||
# dotenv
|
||||
.env
|
||||
|
||||
# virtualenv
|
||||
venv/
|
||||
ENV/
|
||||
|
||||
# Spyder project settings
|
||||
.spyderproject
|
||||
|
||||
# Rope project settings
|
||||
.ropeproject
|
||||
|
||||
|
||||
# Project files
|
||||
.vscode
|
||||
@@ -0,0 +1,201 @@
|
||||
Apache License
|
||||
Version 2.0, January 2004
|
||||
http://www.apache.org/licenses/
|
||||
|
||||
TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
|
||||
|
||||
1. Definitions.
|
||||
|
||||
"License" shall mean the terms and conditions for use, reproduction,
|
||||
and distribution as defined by Sections 1 through 9 of this document.
|
||||
|
||||
"Licensor" shall mean the copyright owner or entity authorized by
|
||||
the copyright owner that is granting the License.
|
||||
|
||||
"Legal Entity" shall mean the union of the acting entity and all
|
||||
other entities that control, are controlled by, or are under common
|
||||
control with that entity. For the purposes of this definition,
|
||||
"control" means (i) the power, direct or indirect, to cause the
|
||||
direction or management of such entity, whether by contract or
|
||||
otherwise, or (ii) ownership of fifty percent (50%) or more of the
|
||||
outstanding shares, or (iii) beneficial ownership of such entity.
|
||||
|
||||
"You" (or "Your") shall mean an individual or Legal Entity
|
||||
exercising permissions granted by this License.
|
||||
|
||||
"Source" form shall mean the preferred form for making modifications,
|
||||
including but not limited to software source code, documentation
|
||||
source, and configuration files.
|
||||
|
||||
"Object" form shall mean any form resulting from mechanical
|
||||
transformation or translation of a Source form, including but
|
||||
not limited to compiled object code, generated documentation,
|
||||
and conversions to other media types.
|
||||
|
||||
"Work" shall mean the work of authorship, whether in Source or
|
||||
Object form, made available under the License, as indicated by a
|
||||
copyright notice that is included in or attached to the work
|
||||
(an example is provided in the Appendix below).
|
||||
|
||||
"Derivative Works" shall mean any work, whether in Source or Object
|
||||
form, that is based on (or derived from) the Work and for which the
|
||||
editorial revisions, annotations, elaborations, or other modifications
|
||||
represent, as a whole, an original work of authorship. For the purposes
|
||||
of this License, Derivative Works shall not include works that remain
|
||||
separable from, or merely link (or bind by name) to the interfaces of,
|
||||
the Work and Derivative Works thereof.
|
||||
|
||||
"Contribution" shall mean any work of authorship, including
|
||||
the original version of the Work and any modifications or additions
|
||||
to that Work or Derivative Works thereof, that is intentionally
|
||||
submitted to Licensor for inclusion in the Work by the copyright owner
|
||||
or by an individual or Legal Entity authorized to submit on behalf of
|
||||
the copyright owner. For the purposes of this definition, "submitted"
|
||||
means any form of electronic, verbal, or written communication sent
|
||||
to the Licensor or its representatives, including but not limited to
|
||||
communication on electronic mailing lists, source code control systems,
|
||||
and issue tracking systems that are managed by, or on behalf of, the
|
||||
Licensor for the purpose of discussing and improving the Work, but
|
||||
excluding communication that is conspicuously marked or otherwise
|
||||
designated in writing by the copyright owner as "Not a Contribution."
|
||||
|
||||
"Contributor" shall mean Licensor and any individual or Legal Entity
|
||||
on behalf of whom a Contribution has been received by Licensor and
|
||||
subsequently incorporated within the Work.
|
||||
|
||||
2. Grant of Copyright License. Subject to the terms and conditions of
|
||||
this License, each Contributor hereby grants to You a perpetual,
|
||||
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
|
||||
copyright license to reproduce, prepare Derivative Works of,
|
||||
publicly display, publicly perform, sublicense, and distribute the
|
||||
Work and such Derivative Works in Source or Object form.
|
||||
|
||||
3. Grant of Patent License. Subject to the terms and conditions of
|
||||
this License, each Contributor hereby grants to You a perpetual,
|
||||
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
|
||||
(except as stated in this section) patent license to make, have made,
|
||||
use, offer to sell, sell, import, and otherwise transfer the Work,
|
||||
where such license applies only to those patent claims licensable
|
||||
by such Contributor that are necessarily infringed by their
|
||||
Contribution(s) alone or by combination of their Contribution(s)
|
||||
with the Work to which such Contribution(s) was submitted. If You
|
||||
institute patent litigation against any entity (including a
|
||||
cross-claim or counterclaim in a lawsuit) alleging that the Work
|
||||
or a Contribution incorporated within the Work constitutes direct
|
||||
or contributory patent infringement, then any patent licenses
|
||||
granted to You under this License for that Work shall terminate
|
||||
as of the date such litigation is filed.
|
||||
|
||||
4. Redistribution. You may reproduce and distribute copies of the
|
||||
Work or Derivative Works thereof in any medium, with or without
|
||||
modifications, and in Source or Object form, provided that You
|
||||
meet the following conditions:
|
||||
|
||||
(a) You must give any other recipients of the Work or
|
||||
Derivative Works a copy of this License; and
|
||||
|
||||
(b) You must cause any modified files to carry prominent notices
|
||||
stating that You changed the files; and
|
||||
|
||||
(c) You must retain, in the Source form of any Derivative Works
|
||||
that You distribute, all copyright, patent, trademark, and
|
||||
attribution notices from the Source form of the Work,
|
||||
excluding those notices that do not pertain to any part of
|
||||
the Derivative Works; and
|
||||
|
||||
(d) If the Work includes a "NOTICE" text file as part of its
|
||||
distribution, then any Derivative Works that You distribute must
|
||||
include a readable copy of the attribution notices contained
|
||||
within such NOTICE file, excluding those notices that do not
|
||||
pertain to any part of the Derivative Works, in at least one
|
||||
of the following places: within a NOTICE text file distributed
|
||||
as part of the Derivative Works; within the Source form or
|
||||
documentation, if provided along with the Derivative Works; or,
|
||||
within a display generated by the Derivative Works, if and
|
||||
wherever such third-party notices normally appear. The contents
|
||||
of the NOTICE file are for informational purposes only and
|
||||
do not modify the License. You may add Your own attribution
|
||||
notices within Derivative Works that You distribute, alongside
|
||||
or as an addendum to the NOTICE text from the Work, provided
|
||||
that such additional attribution notices cannot be construed
|
||||
as modifying the License.
|
||||
|
||||
You may add Your own copyright statement to Your modifications and
|
||||
may provide additional or different license terms and conditions
|
||||
for use, reproduction, or distribution of Your modifications, or
|
||||
for any such Derivative Works as a whole, provided Your use,
|
||||
reproduction, and distribution of the Work otherwise complies with
|
||||
the conditions stated in this License.
|
||||
|
||||
5. Submission of Contributions. Unless You explicitly state otherwise,
|
||||
any Contribution intentionally submitted for inclusion in the Work
|
||||
by You to the Licensor shall be under the terms and conditions of
|
||||
this License, without any additional terms or conditions.
|
||||
Notwithstanding the above, nothing herein shall supersede or modify
|
||||
the terms of any separate license agreement you may have executed
|
||||
with Licensor regarding such Contributions.
|
||||
|
||||
6. Trademarks. This License does not grant permission to use the trade
|
||||
names, trademarks, service marks, or product names of the Licensor,
|
||||
except as required for reasonable and customary use in describing the
|
||||
origin of the Work and reproducing the content of the NOTICE file.
|
||||
|
||||
7. Disclaimer of Warranty. Unless required by applicable law or
|
||||
agreed to in writing, Licensor provides the Work (and each
|
||||
Contributor provides its Contributions) on an "AS IS" BASIS,
|
||||
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
|
||||
implied, including, without limitation, any warranties or conditions
|
||||
of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
|
||||
PARTICULAR PURPOSE. You are solely responsible for determining the
|
||||
appropriateness of using or redistributing the Work and assume any
|
||||
risks associated with Your exercise of permissions under this License.
|
||||
|
||||
8. Limitation of Liability. In no event and under no legal theory,
|
||||
whether in tort (including negligence), contract, or otherwise,
|
||||
unless required by applicable law (such as deliberate and grossly
|
||||
negligent acts) or agreed to in writing, shall any Contributor be
|
||||
liable to You for damages, including any direct, indirect, special,
|
||||
incidental, or consequential damages of any character arising as a
|
||||
result of this License or out of the use or inability to use the
|
||||
Work (including but not limited to damages for loss of goodwill,
|
||||
work stoppage, computer failure or malfunction, or any and all
|
||||
other commercial damages or losses), even if such Contributor
|
||||
has been advised of the possibility of such damages.
|
||||
|
||||
9. Accepting Warranty or Additional Liability. While redistributing
|
||||
the Work or Derivative Works thereof, You may choose to offer,
|
||||
and charge a fee for, acceptance of support, warranty, indemnity,
|
||||
or other liability obligations and/or rights consistent with this
|
||||
License. However, in accepting such obligations, You may act only
|
||||
on Your own behalf and on Your sole responsibility, not on behalf
|
||||
of any other Contributor, and only if You agree to indemnify,
|
||||
defend, and hold each Contributor harmless for any liability
|
||||
incurred by, or claims asserted against, such Contributor by reason
|
||||
of your accepting any such warranty or additional liability.
|
||||
|
||||
END OF TERMS AND CONDITIONS
|
||||
|
||||
APPENDIX: How to apply the Apache License to your work.
|
||||
|
||||
To apply the Apache License to your work, attach the following
|
||||
boilerplate notice, with the fields enclosed by brackets "[]"
|
||||
replaced with your own identifying information. (Don't include
|
||||
the brackets!) The text should be enclosed in the appropriate
|
||||
comment syntax for the file format. We also recommend that a
|
||||
file or class name and description of purpose be included on the
|
||||
same "printed page" as the copyright notice for easier
|
||||
identification within third-party archives.
|
||||
|
||||
Copyright [yyyy] [name of copyright owner]
|
||||
|
||||
Licensed under the Apache License, Version 2.0 (the "License");
|
||||
you may not use this file except in compliance with the License.
|
||||
You may obtain a copy of the License at
|
||||
|
||||
http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
Unless required by applicable law or agreed to in writing, software
|
||||
distributed under the License is distributed on an "AS IS" BASIS,
|
||||
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
See the License for the specific language governing permissions and
|
||||
limitations under the License.
|
||||
@@ -0,0 +1,129 @@
|
||||
<h2 align="center">MODNet: Trimap-Free Portrait Matting in Real Time</h2>
|
||||
|
||||
<div align="center"><i>MODNet: Real-Time Trimap-Free Portrait Matting via Objective Decomposition (AAAI 2022)</i></div>
|
||||
|
||||
<br />
|
||||
|
||||
<img src="doc/gif/homepage_demo.gif" width="100%">
|
||||
|
||||
<div align="center">MODNet is a model for <b>real-time</b> portrait matting with <b>only RGB image input</b></div>
|
||||
<div align="center">MODNet是一个<b>仅需RGB图片输入</b>的<b>实时</b>人像抠图模型</div>
|
||||
|
||||
<br />
|
||||
|
||||
<p align="center">
|
||||
<a href="#online-application-在线应用">Online Application (在线应用)</a> |
|
||||
<a href="#research-demo">Research Demo</a> |
|
||||
<a href="https://arxiv.org/pdf/2011.11961.pdf">AAAI 2022 Paper</a> |
|
||||
<a href="https://youtu.be/PqJ3BRHX3Lc">Supplementary Video</a>
|
||||
</p>
|
||||
|
||||
<p align="center">
|
||||
<a href="#community">Community</a> |
|
||||
<a href="#code">Code</a> |
|
||||
<a href="#ppm-benchmark">PPM Benchmark</a> |
|
||||
<a href="#license">License</a> |
|
||||
<a href="#acknowledgement">Acknowledgement</a> |
|
||||
<a href="#citation">Citation</a> |
|
||||
<a href="#contact">Contact</a>
|
||||
</p>
|
||||
|
||||
---
|
||||
|
||||
|
||||
## Online Application (在线应用)
|
||||
|
||||
The model used in the online demo (unpublished) is only **7M**! Process **2K** resolution image with a **Fast** speed on common PCs or Mobiles! **Beter** than research demos!
|
||||
Please try online portrait image matting on [my personal homepage](https://zhke.io/#/?modnet_demo) for fun!
|
||||
|
||||
在线应用中使用的模型（未发布）大小仅为**7M**！可以在普通PC或移动设备上**快速**处理具有**2K**分辨率的图像！效果比研究示例**更好**！
|
||||
请通过[我的主页](https://zhke.io/#/?modnet_demo)在线尝试图片抠像！
|
||||
|
||||
|
||||
## Research Demo
|
||||
|
||||
All the models behind the following demos are trained on the datasets mentioned in [our paper](https://arxiv.org/pdf/2011.11961.pdf).
|
||||
|
||||
### Portrait Image Matting
|
||||
We provide an [online Colab demo](https://colab.research.google.com/drive/1GANpbKT06aEFiW-Ssx0DQnnEADcXwQG6?usp=sharing) for portrait image matting.
|
||||
It allows you to upload portrait images and predict/visualize/download the alpha mattes.
|
||||
|
||||
<!-- <img src="doc/gif/image_matting_demo.gif" width='40%'> -->
|
||||
|
||||
### Portrait Video Matting
|
||||
We provide two real-time portrait video matting demos based on WebCam. When using the demo, you can move the WebCam around at will.
|
||||
If you have an Ubuntu system, we recommend you to try the [offline demo](demo/video_matting/webcam) to get a higher *fps*. Otherwise, you can access the [online Colab demo](https://colab.research.google.com/drive/1Pt3KDSc2q7WxFvekCnCLD8P0gBEbxm6J?usp=sharing).
|
||||
We also provide an [offline demo](demo/video_matting/custom) that allows you to process custom videos.
|
||||
|
||||
<!-- <img src="doc/gif/video_matting_demo.gif" width='60%'> -->
|
||||
|
||||
|
||||
## Community
|
||||
|
||||
We share some cool applications/extentions of MODNet built by the community.
|
||||
|
||||
<!-- - **WebGUI for Portrait Image Matting** -->
|
||||
<!-- You can try [this WebGUI](https://www.gradio.app/hub/aliabd/modnet) (hosted on [Gradio](https://www.gradio.app/)) for portrait image matting from your browser without code! -->
|
||||
|
||||
- **Colab Demo of Bokeh (Blur Background)**
|
||||
You can try [this Colab demo](https://colab.research.google.com/github/eyaler/avatars4all/blob/master/yarok.ipynb) (built by [@eyaler](https://github.com/eyaler)) to blur the backgroud based on MODNet!
|
||||
|
||||
- **ONNX Version of MODNet**
|
||||
You can convert the pre-trained MODNet to an ONNX model by using [this code](onnx) (provided by [@manthan3C273](https://github.com/manthan3C273)). You can also try [this Colab demo](https://colab.research.google.com/drive/1P3cWtg8fnmu9karZHYDAtmm1vj1rgA-f?usp=sharing) for MODNet image matting (ONNX version).
|
||||
|
||||
- **TorchScript Version of MODNet**
|
||||
You can convert the pre-trained MODNet to an TorchScript model by using [this code](torchscript) (provided by [@yarkable](https://github.com/yarkable)).
|
||||
|
||||
- **TensorRT Version of MODNet**
|
||||
You can access [this Github repository](https://github.com/jkjung-avt/tensorrt_demos) to try the TensorRT version of MODNet (provided by [@jkjung-avt](https://github.com/jkjung-avt)).
|
||||
|
||||
- **Docker Container for MODnet**
|
||||
You can access [this Github repository](https://github.com/nahidalam/modnet_docker) for a containerized version of MODNet with the Docker environment (provided by [@nahidalam](https://github.com/nahidalam)).
|
||||
|
||||
|
||||
There are some resources about MODNet from the community.
|
||||
- [Video from What's AI YouTube Channel](https://youtu.be/rUo0wuVyefU)
|
||||
- [Article from Louis Bouchard's Blog](https://www.louisbouchard.ai/remove-background/)
|
||||
|
||||
|
||||
## Code
|
||||
We provide the [code](src/trainer.py) of MODNet training iteration, including:
|
||||
- **Supervised Training**: Train MODNet on a labeled matting dataset
|
||||
- **SOC Adaptation**: Adapt a trained MODNet to an unlabeled dataset
|
||||
|
||||
In code comments, we provide examples for using the functions.
|
||||
|
||||
|
||||
## PPM Benchmark
|
||||
The PPM benchmark is released in a separate repository [PPM](https://github.com/ZHKKKe/PPM).
|
||||
|
||||
|
||||
## License
|
||||
The code, models, and demos in this repository (excluding GIF files under the folder `doc/gif`) are released under the [Apache License 2.0](https://www.apache.org/licenses/LICENSE-2.0) license.
|
||||
|
||||
|
||||
## Acknowledgement
|
||||
- We thank
|
||||
[@yzhou0919](https://github.com/yzhou0919), [@eyaler](https://github.com/eyaler), [@manthan3C273](https://github.com/manthan3C273), [@yarkable](https://github.com/yarkable), [@jkjung-avt](https://github.com/jkjung-avt), [@manzke](https://github.com/manzke), [@nahidalam](https://github.com/nahidalam),
|
||||
[the Gradio team](https://github.com/gradio-app/gradio), [What's AI YouTube Channel](https://www.youtube.com/channel/UCUzGQrN-lyyc0BWTYoJM_Sg), [Louis Bouchard's Blog](https://www.louisbouchard.ai),
|
||||
for their contributions to this repository or their cool applications/extentions/resources of MODNet.
|
||||
|
||||
|
||||
## Citation
|
||||
If this work helps your research, please consider to cite:
|
||||
|
||||
```bibtex
|
||||
@InProceedings{MODNet,
|
||||
author = {Zhanghan Ke and Jiayu Sun and Kaican Li and Qiong Yan and Rynson W.H. Lau},
|
||||
title = {MODNet: Real-Time Trimap-Free Portrait Matting via Objective Decomposition},
|
||||
booktitle = {AAAI},
|
||||
year = {2022},
|
||||
}
|
||||
```
|
||||
|
||||
|
||||
## Contact
|
||||
This repository is maintained by Zhanghan Ke ([@ZHKKKe](https://github.com/ZHKKKe)).
|
||||
For questions, please contact `kezhanghan@outlook.com`.
|
||||
|
||||
<!-- <img src="doc/gif/commercial_image_matting_model_result.gif" width='100%'> -->
|
||||
Binary file not shown.
|
After Width: | Height: | Size: 382 KiB |
@@ -0,0 +1,2 @@
|
||||
## MODNet - Portrait Image Matting Demo
|
||||
Please try MODNet portrait image matting demo through our [online Colab demo](https://colab.research.google.com/drive/1GANpbKT06aEFiW-Ssx0DQnnEADcXwQG6?usp=sharing).
|
||||
@@ -0,0 +1,105 @@
|
||||
import os
|
||||
import sys
|
||||
import argparse
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
import torchvision.transforms as transforms
|
||||
|
||||
from src.models.modnet import MODNet
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
# define cmd arguments
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument('--input-path', type=str, help='path of input images')
|
||||
parser.add_argument('--output-path', type=str, help='path of output images')
|
||||
parser.add_argument('--ckpt-path', type=str, help='path of pre-trained MODNet')
|
||||
args = parser.parse_args()
|
||||
|
||||
# check input arguments
|
||||
if not os.path.exists(args.input_path):
|
||||
print('Cannot find input path: {0}'.format(args.input_path))
|
||||
exit()
|
||||
if not os.path.exists(args.output_path):
|
||||
print('Cannot find output path: {0}'.format(args.output_path))
|
||||
exit()
|
||||
if not os.path.exists(args.ckpt_path):
|
||||
print('Cannot find ckpt path: {0}'.format(args.ckpt_path))
|
||||
exit()
|
||||
|
||||
# define hyper-parameters
|
||||
ref_size = 512
|
||||
|
||||
# define image to tensor transform
|
||||
im_transform = transforms.Compose(
|
||||
[
|
||||
transforms.ToTensor(),
|
||||
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
|
||||
]
|
||||
)
|
||||
|
||||
# create MODNet and load the pre-trained ckpt
|
||||
modnet = MODNet(backbone_pretrained=False)
|
||||
modnet = nn.DataParallel(modnet)
|
||||
|
||||
if torch.cuda.is_available():
|
||||
modnet = modnet.cuda()
|
||||
weights = torch.load(args.ckpt_path)
|
||||
else:
|
||||
weights = torch.load(args.ckpt_path, map_location=torch.device('cpu'))
|
||||
modnet.load_state_dict(weights)
|
||||
modnet.eval()
|
||||
|
||||
# inference images
|
||||
im_names = os.listdir(args.input_path)
|
||||
for im_name in im_names:
|
||||
print('Process image: {0}'.format(im_name))
|
||||
|
||||
# read image
|
||||
im = Image.open(os.path.join(args.input_path, im_name))
|
||||
|
||||
# unify image channels to 3
|
||||
im = np.asarray(im)
|
||||
if len(im.shape) == 2:
|
||||
im = im[:, :, None]
|
||||
if im.shape[2] == 1:
|
||||
im = np.repeat(im, 3, axis=2)
|
||||
elif im.shape[2] == 4:
|
||||
im = im[:, :, 0:3]
|
||||
|
||||
# convert image to PyTorch tensor
|
||||
im = Image.fromarray(im)
|
||||
im = im_transform(im)
|
||||
|
||||
# add mini-batch dim
|
||||
im = im[None, :, :, :]
|
||||
|
||||
# resize image for input
|
||||
im_b, im_c, im_h, im_w = im.shape
|
||||
if max(im_h, im_w) < ref_size or min(im_h, im_w) > ref_size:
|
||||
if im_w >= im_h:
|
||||
im_rh = ref_size
|
||||
im_rw = int(im_w / im_h * ref_size)
|
||||
elif im_w < im_h:
|
||||
im_rw = ref_size
|
||||
im_rh = int(im_h / im_w * ref_size)
|
||||
else:
|
||||
im_rh = im_h
|
||||
im_rw = im_w
|
||||
|
||||
im_rw = im_rw - im_rw % 32
|
||||
im_rh = im_rh - im_rh % 32
|
||||
im = F.interpolate(im, size=(im_rh, im_rw), mode='area')
|
||||
|
||||
# inference
|
||||
_, _, matte = modnet(im.cuda() if torch.cuda.is_available() else im, True)
|
||||
|
||||
# resize and save matte
|
||||
matte = F.interpolate(matte, size=(im_h, im_w), mode='area')
|
||||
matte = matte[0][0].data.cpu().numpy()
|
||||
matte_name = im_name.split('.')[0] + '.png'
|
||||
Image.fromarray(((matte * 255).astype('uint8')), mode='L').save(os.path.join(args.output_path, matte_name))
|
||||
@@ -0,0 +1,50 @@
|
||||
## MODNet - Custom Portrait Video Matting Demo
|
||||
This is a MODNet portrait video matting demo that allows you to process custom videos.
|
||||
|
||||
### 1. Requirements
|
||||
The basic requirements for this demo are:
|
||||
- Ubuntu System
|
||||
- Python 3+
|
||||
|
||||
|
||||
### 2. Introduction
|
||||
We use ~400 unlabeled video clips (divided into ~50,000 frames) downloaded from the internet to perform SOC to adapt MODNet to the video domain. **Nonetheless, due to insufficient labeled training data (~3k labeled foregrounds), our model may still make errors in portrait semantics estimation under challenging scenes.** Besides, this demo does not currently support the OFD trick.
|
||||
|
||||
|
||||
For a better experience, please make sure your videos satisfy:
|
||||
|
||||
* the portrait and background are distinguishable, <i>i.e.</i>, are not similar
|
||||
* captured in soft and bright ambient lighting
|
||||
* the contents do not move too fast
|
||||
|
||||
### 3. Run Demo
|
||||
We recommend creating a new conda virtual environment to run this demo, as follow:
|
||||
|
||||
1. Clone the MODNet repository:
|
||||
```
|
||||
git clone https://github.com/ZHKKKe/MODNet.git
|
||||
cd MODNet
|
||||
```
|
||||
|
||||
2. Download the pre-trained model from this [link](https://drive.google.com/file/d/1Nf1ZxeJZJL8Qx9KadcYYyEmmlKhTADxX/view?usp=sharing) and put it into the folder `MODNet/pretrained/`.
|
||||
|
||||
|
||||
3. Create a conda virtual environment named `modnet` (if it doesn't exist) and activate it. Here we use `python=3.6` as an example:
|
||||
```
|
||||
conda create -n modnet python=3.6
|
||||
source activate modnet
|
||||
```
|
||||
|
||||
4. Install the required python dependencies (please make sure your CUDA version is supported by the PyTorch version installed):
|
||||
```
|
||||
pip install -r demo/video_matting/custom/requirements.txt
|
||||
```
|
||||
|
||||
5. Execute the main code:
|
||||
```
|
||||
python -m demo.video_matting.custom.run --video YOUR_VIDEO_PATH
|
||||
```
|
||||
where `YOUR_VIDEO_PATH` is the specific path of your video.
|
||||
There are some optional arguments:
|
||||
- `--result-type (default=fg)` : fg - save the alpha matte; fg - save the foreground
|
||||
- `--fps (default=30)` : fps of the result video
|
||||
@@ -0,0 +1,6 @@
|
||||
numpy
|
||||
Pillow
|
||||
opencv-python
|
||||
torch >= 1.0.0
|
||||
torchvision
|
||||
tqdm
|
||||
@@ -0,0 +1,114 @@
|
||||
import os
|
||||
import cv2
|
||||
import argparse
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
from tqdm import tqdm
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torchvision.transforms as transforms
|
||||
|
||||
from src.models.modnet import MODNet
|
||||
|
||||
|
||||
torch_transforms = transforms.Compose(
|
||||
[
|
||||
transforms.ToTensor(),
|
||||
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
|
||||
]
|
||||
)
|
||||
|
||||
|
||||
def matting(video, result, alpha_matte=False, fps=30):
|
||||
# video capture
|
||||
vc = cv2.VideoCapture(video)
|
||||
|
||||
if vc.isOpened():
|
||||
rval, frame = vc.read()
|
||||
else:
|
||||
rval = False
|
||||
|
||||
if not rval:
|
||||
print('Failed to read the video: {0}'.format(video))
|
||||
exit()
|
||||
|
||||
num_frame = vc.get(cv2.CAP_PROP_FRAME_COUNT)
|
||||
h, w = frame.shape[:2]
|
||||
if w >= h:
|
||||
rh = 512
|
||||
rw = int(w / h * 512)
|
||||
else:
|
||||
rw = 512
|
||||
rh = int(h / w * 512)
|
||||
rh = rh - rh % 32
|
||||
rw = rw - rw % 32
|
||||
|
||||
# video writer
|
||||
fourcc = cv2.VideoWriter_fourcc(*'mp4v')
|
||||
video_writer = cv2.VideoWriter(result, fourcc, fps, (w, h))
|
||||
|
||||
print('Start matting...')
|
||||
with tqdm(range(int(num_frame)))as t:
|
||||
for c in t:
|
||||
frame_np = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
|
||||
frame_np = cv2.resize(frame_np, (rw, rh), cv2.INTER_AREA)
|
||||
|
||||
frame_PIL = Image.fromarray(frame_np)
|
||||
frame_tensor = torch_transforms(frame_PIL)
|
||||
frame_tensor = frame_tensor[None, :, :, :]
|
||||
if GPU:
|
||||
frame_tensor = frame_tensor.cuda()
|
||||
|
||||
with torch.no_grad():
|
||||
_, _, matte_tensor = modnet(frame_tensor, True)
|
||||
|
||||
matte_tensor = matte_tensor.repeat(1, 3, 1, 1)
|
||||
matte_np = matte_tensor[0].data.cpu().numpy().transpose(1, 2, 0)
|
||||
if alpha_matte:
|
||||
view_np = matte_np * np.full(frame_np.shape, 255.0)
|
||||
else:
|
||||
view_np = matte_np * frame_np + (1 - matte_np) * np.full(frame_np.shape, 255.0)
|
||||
view_np = cv2.cvtColor(view_np.astype(np.uint8), cv2.COLOR_RGB2BGR)
|
||||
view_np = cv2.resize(view_np, (w, h))
|
||||
video_writer.write(view_np)
|
||||
|
||||
rval, frame = vc.read()
|
||||
c += 1
|
||||
|
||||
video_writer.release()
|
||||
print('Save the result video to {0}'.format(result))
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument('--video', type=str, required=True, help='input video file')
|
||||
parser.add_argument('--result-type', type=str, default='fg', choices=['fg', 'matte'],
|
||||
help='matte - save the alpha matte; fg - save the foreground')
|
||||
parser.add_argument('--fps', type=int, default=30, help='fps of the result video')
|
||||
|
||||
print('Get CMD Arguments...')
|
||||
args = parser.parse_args()
|
||||
|
||||
if not os.path.exists(args.video):
|
||||
print('Cannot find the input video: {0}'.format(args.video))
|
||||
exit()
|
||||
|
||||
print('Load pre-trained MODNet...')
|
||||
pretrained_ckpt = './pretrained/modnet_webcam_portrait_matting.ckpt'
|
||||
modnet = MODNet(backbone_pretrained=False)
|
||||
modnet = nn.DataParallel(modnet)
|
||||
|
||||
GPU = True if torch.cuda.device_count() > 0 else False
|
||||
if GPU:
|
||||
print('Use GPU...')
|
||||
modnet = modnet.cuda()
|
||||
modnet.load_state_dict(torch.load(pretrained_ckpt))
|
||||
else:
|
||||
print('Use CPU...')
|
||||
modnet.load_state_dict(torch.load(pretrained_ckpt, map_location=torch.device('cpu')))
|
||||
modnet.eval()
|
||||
|
||||
result = os.path.splitext(args.video)[0] + '_{0}.mp4'.format(args.result_type)
|
||||
alpha_matte = True if args.result_type == 'matte' else False
|
||||
matting(args.video, result, alpha_matte, args.fps)
|
||||
@@ -0,0 +1,52 @@
|
||||
## MODNet - WebCam-Based Portrait Video Matting Demo
|
||||
This is a MODNet portrait video matting demo based on WebCam. It will call your local WebCam and display the matting results in real time. The demo can run under CPU or GPU.
|
||||
|
||||
### 1. Requirements
|
||||
The basic requirements for this demo are:
|
||||
- Ubuntu System
|
||||
- WebCam
|
||||
- Python 3+
|
||||
|
||||
**NOTE**: If your device does not satisfy the above conditions, please try our [online Colab demo](https://colab.research.google.com/drive/1Pt3KDSc2q7WxFvekCnCLD8P0gBEbxm6J?usp=sharing).
|
||||
|
||||
|
||||
### 2. Introduction
|
||||
We use ~400 unlabeled video clips (divided into ~50,000 frames) downloaded from the internet to perform SOC to adapt MODNet to the video domain. **Nonetheless, due to insufficient labeled training data (~3k labeled foregrounds), our model may still make errors in portrait semantics estimation under challenging scenes.** Besides, this demo does not currently support the OFD trick, which will be provided soon.
|
||||
|
||||
For a better experience, please:
|
||||
|
||||
* make sure the portrait and background are distinguishable, <i>i.e.</i>, are not similar
|
||||
* run in soft and bright ambient lighting
|
||||
* do not be too close or too far from the WebCam
|
||||
* do not move too fast
|
||||
|
||||
### 3. Run Demo
|
||||
We recommend creating a new conda virtual environment to run this demo, as follow:
|
||||
|
||||
1. Clone the MODNet repository:
|
||||
```
|
||||
git clone https://github.com/ZHKKKe/MODNet.git
|
||||
cd MODNet
|
||||
```
|
||||
|
||||
2. Download the pre-trained model from this [link](https://drive.google.com/file/d/1Nf1ZxeJZJL8Qx9KadcYYyEmmlKhTADxX/view?usp=sharing) and put it into the folder `MODNet/pretrained/`.
|
||||
|
||||
|
||||
3. Create a conda virtual environment named `modnet` (if it doesn't exist) and activate it. Here we use `python=3.6` as an example:
|
||||
```
|
||||
conda create -n modnet python=3.6
|
||||
source activate modnet
|
||||
```
|
||||
|
||||
4. Install the required python dependencies (please make sure your CUDA version is supported by the PyTorch version installed):
|
||||
```
|
||||
pip install -r demo/video_matting/webcam/requirements.txt
|
||||
```
|
||||
|
||||
5. Execute the main code:
|
||||
```
|
||||
python -m demo.video_matting.webcam.run
|
||||
```
|
||||
|
||||
### 4. Acknowledgement
|
||||
We thank [@tkianai](https://github.com/tkianai) and [@mazhar004](https://github.com/mazhar004) for their contributions to making this demo available for CPU use.
|
||||
@@ -0,0 +1,5 @@
|
||||
numpy
|
||||
Pillow
|
||||
opencv-python
|
||||
torch >= 1.0.0
|
||||
torchvision
|
||||
@@ -0,0 +1,67 @@
|
||||
import cv2
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torchvision.transforms as transforms
|
||||
|
||||
from src.models.modnet import MODNet
|
||||
|
||||
|
||||
torch_transforms = transforms.Compose(
|
||||
[
|
||||
transforms.ToTensor(),
|
||||
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
|
||||
]
|
||||
)
|
||||
|
||||
print('Load pre-trained MODNet...')
|
||||
pretrained_ckpt = './pretrained/modnet_webcam_portrait_matting.ckpt'
|
||||
modnet = MODNet(backbone_pretrained=False)
|
||||
modnet = nn.DataParallel(modnet)
|
||||
|
||||
GPU = True if torch.cuda.device_count() > 0 else False
|
||||
if GPU:
|
||||
print('Use GPU...')
|
||||
modnet = modnet.cuda()
|
||||
modnet.load_state_dict(torch.load(pretrained_ckpt))
|
||||
else:
|
||||
print('Use CPU...')
|
||||
modnet.load_state_dict(torch.load(pretrained_ckpt, map_location=torch.device('cpu')))
|
||||
|
||||
modnet.eval()
|
||||
|
||||
print('Init WebCam...')
|
||||
cap = cv2.VideoCapture(0)
|
||||
cap.set(cv2.CAP_PROP_FRAME_WIDTH, 1280)
|
||||
cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 720)
|
||||
|
||||
print('Start matting...')
|
||||
while(True):
|
||||
_, frame_np = cap.read()
|
||||
frame_np = cv2.cvtColor(frame_np, cv2.COLOR_BGR2RGB)
|
||||
frame_np = cv2.resize(frame_np, (910, 512), cv2.INTER_AREA)
|
||||
frame_np = frame_np[:, 120:792, :]
|
||||
frame_np = cv2.flip(frame_np, 1)
|
||||
|
||||
frame_PIL = Image.fromarray(frame_np)
|
||||
frame_tensor = torch_transforms(frame_PIL)
|
||||
frame_tensor = frame_tensor[None, :, :, :]
|
||||
if GPU:
|
||||
frame_tensor = frame_tensor.cuda()
|
||||
|
||||
with torch.no_grad():
|
||||
_, _, matte_tensor = modnet(frame_tensor, True)
|
||||
|
||||
matte_tensor = matte_tensor.repeat(1, 3, 1, 1)
|
||||
matte_np = matte_tensor[0].data.cpu().numpy().transpose(1, 2, 0)
|
||||
fg_np = matte_np * frame_np + (1 - matte_np) * np.full(frame_np.shape, 255.0)
|
||||
view_np = np.uint8(np.concatenate((frame_np, fg_np), axis=1))
|
||||
view_np = cv2.cvtColor(view_np, cv2.COLOR_RGB2BGR)
|
||||
|
||||
cv2.imshow('MODNet - WebCam [Press \'Q\' To Exit]', view_np)
|
||||
if cv2.waitKey(1) & 0xFF == ord('q'):
|
||||
break
|
||||
|
||||
print('Exit...')
|
||||
@@ -0,0 +1,30 @@
|
||||
## MODNet - ONNX Model
|
||||
|
||||
This ONNX version of MODNet is provided by [@manthan3C273](https://github.com/manthan3C273) from the community.
|
||||
Please note that the PyTorch version required for this ONNX export function is higher than the official MODNet code (torch==1.7.1 is recommended).
|
||||
|
||||
You can try **MODNet - Image Matting Demo (ONNX version)** in [this Colab](https://colab.research.google.com/drive/1P3cWtg8fnmu9karZHYDAtmm1vj1rgA-f?usp=sharing).
|
||||
You can also download the ONNX version of the official **Image Matting Model** from [this link](https://drive.google.com/file/d/1cgycTQlYXpTh26gB9FTnthE7AvruV8hd/view?usp=sharing).
|
||||
|
||||
To export the ONNX version of MODNet (assuming you are currently in project root directory):
|
||||
1. Download the pre-trained **Image Matting Model** from this [link](https://drive.google.com/drive/folders/1umYmlCulvIFNaqPjwod1SayFmSRHziyR?usp=sharing) and put the model into the folder `MODNet/pretrained/`.
|
||||
|
||||
2. Install all dependencies by:
|
||||
```
|
||||
pip install -r onnx/requirements.txt
|
||||
```
|
||||
|
||||
3. Export the ONNX version of MODNet by:
|
||||
```shell
|
||||
python -m onnx.export_onnx \
|
||||
--ckpt-path=pretrained/modnet_photographic_portrait_matting.ckpt \
|
||||
--output-path=pretrained/modnet_photographic_portrait_matting.onnx
|
||||
```
|
||||
|
||||
4. Inference the ONNX model by:
|
||||
```shell
|
||||
python -m onnx.inference_onnx \
|
||||
--image-path=$FILENAME_OF_INPUT_IMAGE$ \
|
||||
--output-path=$FILENAME_OF_OUTPUT_MATTE$ \
|
||||
--model-path=pretrained/modnet_photographic_portrait_matting.onnx
|
||||
```
|
||||
@@ -0,0 +1,55 @@
|
||||
"""
|
||||
Export ONNX model of MODNet with:
|
||||
input shape: (batch_size, 3, height, width)
|
||||
output shape: (batch_size, 1, height, width)
|
||||
|
||||
Arguments:
|
||||
--ckpt-path: path of the checkpoint that will be converted
|
||||
--output-path: path for saving the ONNX model
|
||||
|
||||
Example:
|
||||
python export_onnx.py \
|
||||
--ckpt-path=modnet_photographic_portrait_matting.ckpt \
|
||||
--output-path=modnet_photographic_portrait_matting.onnx
|
||||
"""
|
||||
|
||||
import os
|
||||
import argparse
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from torch.autograd import Variable
|
||||
|
||||
from . import modnet_onnx
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
# define cmd arguments
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument('--ckpt-path', type=str, required=True, help='path of the checkpoint that will be converted')
|
||||
parser.add_argument('--output-path', type=str, required=True, help='path for saving the ONNX model')
|
||||
args = parser.parse_args()
|
||||
|
||||
# check input arguments
|
||||
if not os.path.exists(args.ckpt_path):
|
||||
print('Cannot find checkpoint path: {0}'.format(args.ckpt_path))
|
||||
exit()
|
||||
|
||||
# define model & load checkpoint
|
||||
modnet = modnet_onnx.MODNet(backbone_pretrained=False)
|
||||
modnet = nn.DataParallel(modnet).cuda()
|
||||
state_dict = torch.load(args.ckpt_path)
|
||||
modnet.load_state_dict(state_dict)
|
||||
modnet.eval()
|
||||
|
||||
# prepare dummy_input
|
||||
batch_size = 1
|
||||
height = 512
|
||||
width = 512
|
||||
dummy_input = Variable(torch.randn(batch_size, 3, height, width)).cuda()
|
||||
|
||||
# export to onnx model
|
||||
torch.onnx.export(
|
||||
modnet.module, dummy_input, args.output_path, export_params = True,
|
||||
input_names = ['input'], output_names = ['output'],
|
||||
dynamic_axes = {'input': {0:'batch_size', 2:'height', 3:'width'}, 'output': {0: 'batch_size', 2: 'height', 3: 'width'}})
|
||||
@@ -0,0 +1,104 @@
|
||||
"""
|
||||
Inference ONNX model of MODNet
|
||||
|
||||
Arguments:
|
||||
--image-path: path of the input image (a file)
|
||||
--output-path: path for saving the predicted alpha matte (a file)
|
||||
--model-path: path of the ONNX model
|
||||
|
||||
Example:
|
||||
python inference_onnx.py \
|
||||
--image-path=demo.jpg --output-path=matte.png --model-path=modnet.onnx
|
||||
"""
|
||||
|
||||
import os
|
||||
import cv2
|
||||
import argparse
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
|
||||
import onnx
|
||||
import onnxruntime
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
# define cmd arguments
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument('--image-path', type=str, help='path of the input image (a file)')
|
||||
parser.add_argument('--output-path', type=str, help='paht for saving the predicted alpha matte (a file)')
|
||||
parser.add_argument('--model-path', type=str, help='path of the ONNX model')
|
||||
args = parser.parse_args()
|
||||
|
||||
# check input arguments
|
||||
if not os.path.exists(args.image_path):
|
||||
print('Cannot find the input image: {0}'.format(args.image_path))
|
||||
exit()
|
||||
if not os.path.exists(args.model_path):
|
||||
print('Cannot find the ONXX model: {0}'.format(args.model_path))
|
||||
exit()
|
||||
|
||||
ref_size = 512
|
||||
|
||||
# Get x_scale_factor & y_scale_factor to resize image
|
||||
def get_scale_factor(im_h, im_w, ref_size):
|
||||
|
||||
if max(im_h, im_w) < ref_size or min(im_h, im_w) > ref_size:
|
||||
if im_w >= im_h:
|
||||
im_rh = ref_size
|
||||
im_rw = int(im_w / im_h * ref_size)
|
||||
elif im_w < im_h:
|
||||
im_rw = ref_size
|
||||
im_rh = int(im_h / im_w * ref_size)
|
||||
else:
|
||||
im_rh = im_h
|
||||
im_rw = im_w
|
||||
|
||||
im_rw = im_rw - im_rw % 32
|
||||
im_rh = im_rh - im_rh % 32
|
||||
|
||||
x_scale_factor = im_rw / im_w
|
||||
y_scale_factor = im_rh / im_h
|
||||
|
||||
return x_scale_factor, y_scale_factor
|
||||
|
||||
##############################################
|
||||
# Main Inference part
|
||||
##############################################
|
||||
|
||||
# read image
|
||||
im = cv2.imread(args.image_path)
|
||||
im = cv2.cvtColor(im, cv2.COLOR_BGR2RGB)
|
||||
|
||||
# unify image channels to 3
|
||||
if len(im.shape) == 2:
|
||||
im = im[:, :, None]
|
||||
if im.shape[2] == 1:
|
||||
im = np.repeat(im, 3, axis=2)
|
||||
elif im.shape[2] == 4:
|
||||
im = im[:, :, 0:3]
|
||||
|
||||
# normalize values to scale it between -1 to 1
|
||||
im = (im - 127.5) / 127.5
|
||||
|
||||
im_h, im_w, im_c = im.shape
|
||||
x, y = get_scale_factor(im_h, im_w, ref_size)
|
||||
|
||||
# resize image
|
||||
im = cv2.resize(im, None, fx = x, fy = y, interpolation = cv2.INTER_AREA)
|
||||
|
||||
# prepare input shape
|
||||
im = np.transpose(im)
|
||||
im = np.swapaxes(im, 1, 2)
|
||||
im = np.expand_dims(im, axis = 0).astype('float32')
|
||||
|
||||
# Initialize session and get prediction
|
||||
session = onnxruntime.InferenceSession(args.model_path, None)
|
||||
input_name = session.get_inputs()[0].name
|
||||
output_name = session.get_outputs()[0].name
|
||||
result = session.run([output_name], {input_name: im})
|
||||
|
||||
# refine matte
|
||||
matte = (np.squeeze(result[0]) * 255).astype('uint8')
|
||||
matte = cv2.resize(matte, dsize=(im_w, im_h), interpolation = cv2.INTER_AREA)
|
||||
|
||||
cv2.imwrite(args.output_path, matte)
|
||||
@@ -0,0 +1,252 @@
|
||||
"""
|
||||
This file contains a modified version of the original file `modnet.py` without
|
||||
`pred_semantic` and `pred_details` as these both returns None when `inference=True`
|
||||
|
||||
And it does not contain `inference` argument which will make it easier to
|
||||
convert checkpoint to ONNX model.
|
||||
"""
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
from src.models.backbones import SUPPORTED_BACKBONES
|
||||
|
||||
|
||||
#------------------------------------------------------------------------------
|
||||
# MODNet Basic Modules
|
||||
#------------------------------------------------------------------------------
|
||||
|
||||
class IBNorm(nn.Module):
|
||||
""" Combine Instance Norm and Batch Norm into One Layer
|
||||
"""
|
||||
|
||||
def __init__(self, in_channels):
|
||||
super(IBNorm, self).__init__()
|
||||
in_channels = in_channels
|
||||
self.bnorm_channels = int(in_channels / 2)
|
||||
self.inorm_channels = in_channels - self.bnorm_channels
|
||||
|
||||
self.bnorm = nn.BatchNorm2d(self.bnorm_channels, affine=True)
|
||||
self.inorm = nn.InstanceNorm2d(self.inorm_channels, affine=False)
|
||||
|
||||
def forward(self, x):
|
||||
bn_x = self.bnorm(x[:, :self.bnorm_channels, ...].contiguous())
|
||||
in_x = self.inorm(x[:, self.bnorm_channels:, ...].contiguous())
|
||||
|
||||
return torch.cat((bn_x, in_x), 1)
|
||||
|
||||
|
||||
class Conv2dIBNormRelu(nn.Module):
|
||||
""" Convolution + IBNorm + ReLu
|
||||
"""
|
||||
|
||||
def __init__(self, in_channels, out_channels, kernel_size,
|
||||
stride=1, padding=0, dilation=1, groups=1, bias=True,
|
||||
with_ibn=True, with_relu=True):
|
||||
super(Conv2dIBNormRelu, self).__init__()
|
||||
|
||||
layers = [
|
||||
nn.Conv2d(in_channels, out_channels, kernel_size,
|
||||
stride=stride, padding=padding, dilation=dilation,
|
||||
groups=groups, bias=bias)
|
||||
]
|
||||
|
||||
if with_ibn:
|
||||
layers.append(IBNorm(out_channels))
|
||||
if with_relu:
|
||||
layers.append(nn.ReLU(inplace=True))
|
||||
|
||||
self.layers = nn.Sequential(*layers)
|
||||
|
||||
def forward(self, x):
|
||||
return self.layers(x)
|
||||
|
||||
|
||||
class SEBlock(nn.Module):
|
||||
""" SE Block Proposed in https://arxiv.org/pdf/1709.01507.pdf
|
||||
"""
|
||||
|
||||
def __init__(self, in_channels, out_channels, reduction=1):
|
||||
super(SEBlock, self).__init__()
|
||||
self.pool = nn.AdaptiveAvgPool2d(1)
|
||||
self.fc = nn.Sequential(
|
||||
nn.Linear(in_channels, int(in_channels // reduction), bias=False),
|
||||
nn.ReLU(inplace=True),
|
||||
nn.Linear(int(in_channels // reduction), out_channels, bias=False),
|
||||
nn.Sigmoid()
|
||||
)
|
||||
|
||||
def forward(self, x):
|
||||
b, c, _, _ = x.size()
|
||||
w = self.pool(x).view(b, c)
|
||||
w = self.fc(w).view(b, c, 1, 1)
|
||||
|
||||
return x * w.expand_as(x)
|
||||
|
||||
|
||||
#------------------------------------------------------------------------------
|
||||
# MODNet Branches
|
||||
#------------------------------------------------------------------------------
|
||||
|
||||
class LRBranch(nn.Module):
|
||||
""" Low Resolution Branch of MODNet
|
||||
"""
|
||||
|
||||
def __init__(self, backbone):
|
||||
super(LRBranch, self).__init__()
|
||||
|
||||
enc_channels = backbone.enc_channels
|
||||
|
||||
self.backbone = backbone
|
||||
self.se_block = SEBlock(enc_channels[4], enc_channels[4], reduction=4)
|
||||
self.conv_lr16x = Conv2dIBNormRelu(enc_channels[4], enc_channels[3], 5, stride=1, padding=2)
|
||||
self.conv_lr8x = Conv2dIBNormRelu(enc_channels[3], enc_channels[2], 5, stride=1, padding=2)
|
||||
self.conv_lr = Conv2dIBNormRelu(enc_channels[2], 1, kernel_size=3, stride=2, padding=1, with_ibn=False, with_relu=False)
|
||||
|
||||
def forward(self, img):
|
||||
enc_features = self.backbone.forward(img)
|
||||
enc2x, enc4x, enc32x = enc_features[0], enc_features[1], enc_features[4]
|
||||
|
||||
enc32x = self.se_block(enc32x)
|
||||
lr16x = F.interpolate(enc32x, scale_factor=2, mode='bilinear', align_corners=False)
|
||||
lr16x = self.conv_lr16x(lr16x)
|
||||
lr8x = F.interpolate(lr16x, scale_factor=2, mode='bilinear', align_corners=False)
|
||||
lr8x = self.conv_lr8x(lr8x)
|
||||
|
||||
return lr8x, [enc2x, enc4x]
|
||||
|
||||
|
||||
class HRBranch(nn.Module):
|
||||
""" High Resolution Branch of MODNet
|
||||
"""
|
||||
|
||||
def __init__(self, hr_channels, enc_channels):
|
||||
super(HRBranch, self).__init__()
|
||||
|
||||
self.tohr_enc2x = Conv2dIBNormRelu(enc_channels[0], hr_channels, 1, stride=1, padding=0)
|
||||
self.conv_enc2x = Conv2dIBNormRelu(hr_channels + 3, hr_channels, 3, stride=2, padding=1)
|
||||
|
||||
self.tohr_enc4x = Conv2dIBNormRelu(enc_channels[1], hr_channels, 1, stride=1, padding=0)
|
||||
self.conv_enc4x = Conv2dIBNormRelu(2 * hr_channels, 2 * hr_channels, 3, stride=1, padding=1)
|
||||
|
||||
self.conv_hr4x = nn.Sequential(
|
||||
Conv2dIBNormRelu(3 * hr_channels + 3, 2 * hr_channels, 3, stride=1, padding=1),
|
||||
Conv2dIBNormRelu(2 * hr_channels, 2 * hr_channels, 3, stride=1, padding=1),
|
||||
Conv2dIBNormRelu(2 * hr_channels, hr_channels, 3, stride=1, padding=1),
|
||||
)
|
||||
|
||||
self.conv_hr2x = nn.Sequential(
|
||||
Conv2dIBNormRelu(2 * hr_channels, 2 * hr_channels, 3, stride=1, padding=1),
|
||||
Conv2dIBNormRelu(2 * hr_channels, hr_channels, 3, stride=1, padding=1),
|
||||
Conv2dIBNormRelu(hr_channels, hr_channels, 3, stride=1, padding=1),
|
||||
Conv2dIBNormRelu(hr_channels, hr_channels, 3, stride=1, padding=1),
|
||||
)
|
||||
|
||||
self.conv_hr = nn.Sequential(
|
||||
Conv2dIBNormRelu(hr_channels + 3, hr_channels, 3, stride=1, padding=1),
|
||||
Conv2dIBNormRelu(hr_channels, 1, kernel_size=1, stride=1, padding=0, with_ibn=False, with_relu=False),
|
||||
)
|
||||
|
||||
def forward(self, img, enc2x, enc4x, lr8x):
|
||||
img2x = F.interpolate(img, scale_factor=1/2, mode='bilinear', align_corners=False)
|
||||
img4x = F.interpolate(img, scale_factor=1/4, mode='bilinear', align_corners=False)
|
||||
|
||||
enc2x = self.tohr_enc2x(enc2x)
|
||||
hr4x = self.conv_enc2x(torch.cat((img2x, enc2x), dim=1))
|
||||
|
||||
enc4x = self.tohr_enc4x(enc4x)
|
||||
hr4x = self.conv_enc4x(torch.cat((hr4x, enc4x), dim=1))
|
||||
|
||||
lr4x = F.interpolate(lr8x, scale_factor=2, mode='bilinear', align_corners=False)
|
||||
hr4x = self.conv_hr4x(torch.cat((hr4x, lr4x, img4x), dim=1))
|
||||
|
||||
hr2x = F.interpolate(hr4x, scale_factor=2, mode='bilinear', align_corners=False)
|
||||
hr2x = self.conv_hr2x(torch.cat((hr2x, enc2x), dim=1))
|
||||
|
||||
return hr2x
|
||||
|
||||
|
||||
class FusionBranch(nn.Module):
|
||||
""" Fusion Branch of MODNet
|
||||
"""
|
||||
|
||||
def __init__(self, hr_channels, enc_channels):
|
||||
super(FusionBranch, self).__init__()
|
||||
self.conv_lr4x = Conv2dIBNormRelu(enc_channels[2], hr_channels, 5, stride=1, padding=2)
|
||||
|
||||
self.conv_f2x = Conv2dIBNormRelu(2 * hr_channels, hr_channels, 3, stride=1, padding=1)
|
||||
self.conv_f = nn.Sequential(
|
||||
Conv2dIBNormRelu(hr_channels + 3, int(hr_channels / 2), 3, stride=1, padding=1),
|
||||
Conv2dIBNormRelu(int(hr_channels / 2), 1, 1, stride=1, padding=0, with_ibn=False, with_relu=False),
|
||||
)
|
||||
|
||||
def forward(self, img, lr8x, hr2x):
|
||||
lr4x = F.interpolate(lr8x, scale_factor=2, mode='bilinear', align_corners=False)
|
||||
lr4x = self.conv_lr4x(lr4x)
|
||||
lr2x = F.interpolate(lr4x, scale_factor=2, mode='bilinear', align_corners=False)
|
||||
|
||||
f2x = self.conv_f2x(torch.cat((lr2x, hr2x), dim=1))
|
||||
f = F.interpolate(f2x, scale_factor=2, mode='bilinear', align_corners=False)
|
||||
f = self.conv_f(torch.cat((f, img), dim=1))
|
||||
pred_matte = torch.sigmoid(f)
|
||||
|
||||
return pred_matte
|
||||
|
||||
|
||||
#------------------------------------------------------------------------------
|
||||
# MODNet
|
||||
#------------------------------------------------------------------------------
|
||||
|
||||
class MODNet(nn.Module):
|
||||
""" Architecture of MODNet
|
||||
"""
|
||||
|
||||
def __init__(self, in_channels=3, hr_channels=32, backbone_arch='mobilenetv2', backbone_pretrained=True):
|
||||
super(MODNet, self).__init__()
|
||||
|
||||
self.in_channels = in_channels
|
||||
self.hr_channels = hr_channels
|
||||
self.backbone_arch = backbone_arch
|
||||
self.backbone_pretrained = backbone_pretrained
|
||||
|
||||
self.backbone = SUPPORTED_BACKBONES[self.backbone_arch](self.in_channels)
|
||||
|
||||
self.lr_branch = LRBranch(self.backbone)
|
||||
self.hr_branch = HRBranch(self.hr_channels, self.backbone.enc_channels)
|
||||
self.f_branch = FusionBranch(self.hr_channels, self.backbone.enc_channels)
|
||||
|
||||
for m in self.modules():
|
||||
if isinstance(m, nn.Conv2d):
|
||||
self._init_conv(m)
|
||||
elif isinstance(m, nn.BatchNorm2d) or isinstance(m, nn.InstanceNorm2d):
|
||||
self._init_norm(m)
|
||||
|
||||
if self.backbone_pretrained:
|
||||
self.backbone.load_pretrained_ckpt()
|
||||
|
||||
def forward(self, img):
|
||||
lr8x, [enc2x, enc4x] = self.lr_branch(img)
|
||||
hr2x = self.hr_branch(img, enc2x, enc4x, lr8x)
|
||||
pred_matte = self.f_branch(img, lr8x, hr2x)
|
||||
|
||||
return pred_matte
|
||||
|
||||
def freeze_norm(self):
|
||||
norm_types = [nn.BatchNorm2d, nn.InstanceNorm2d]
|
||||
for m in self.modules():
|
||||
for n in norm_types:
|
||||
if isinstance(m, n):
|
||||
m.eval()
|
||||
continue
|
||||
|
||||
def _init_conv(self, conv):
|
||||
nn.init.kaiming_uniform_(
|
||||
conv.weight, a=0, mode='fan_in', nonlinearity='relu')
|
||||
if conv.bias is not None:
|
||||
nn.init.constant_(conv.bias, 0)
|
||||
|
||||
def _init_norm(self, norm):
|
||||
if norm.weight is not None:
|
||||
nn.init.constant_(norm.weight, 1)
|
||||
nn.init.constant_(norm.bias, 0)
|
||||
@@ -0,0 +1,4 @@
|
||||
onnx==1.8.1
|
||||
onnxruntime==1.6.0
|
||||
opencv-python==4.5.1.48
|
||||
torch==1.7.1
|
||||
@@ -0,0 +1,2 @@
|
||||
## MODNet - Pre-Trained Models
|
||||
This folder is used to save the official pre-trained models of MODNet. You can download them from this [link](https://drive.google.com/drive/folders/1umYmlCulvIFNaqPjwod1SayFmSRHziyR?usp=sharing).
|
||||
@@ -0,0 +1,10 @@
|
||||
from .wrapper import *
|
||||
|
||||
|
||||
#------------------------------------------------------------------------------
|
||||
# Replaceable Backbones
|
||||
#------------------------------------------------------------------------------
|
||||
|
||||
SUPPORTED_BACKBONES = {
|
||||
'mobilenetv2': MobileNetV2Backbone,
|
||||
}
|
||||
@@ -0,0 +1,199 @@
|
||||
""" This file is adapted from https://github.com/thuyngch/Human-Segmentation-PyTorch"""
|
||||
|
||||
import math
|
||||
import json
|
||||
from functools import reduce
|
||||
|
||||
import torch
|
||||
from torch import nn
|
||||
|
||||
|
||||
#------------------------------------------------------------------------------
|
||||
# Useful functions
|
||||
#------------------------------------------------------------------------------
|
||||
|
||||
def _make_divisible(v, divisor, min_value=None):
|
||||
if min_value is None:
|
||||
min_value = divisor
|
||||
new_v = max(min_value, int(v + divisor / 2) // divisor * divisor)
|
||||
# Make sure that round down does not go down by more than 10%.
|
||||
if new_v < 0.9 * v:
|
||||
new_v += divisor
|
||||
return new_v
|
||||
|
||||
|
||||
def conv_bn(inp, oup, stride):
|
||||
return nn.Sequential(
|
||||
nn.Conv2d(inp, oup, 3, stride, 1, bias=False),
|
||||
nn.BatchNorm2d(oup),
|
||||
nn.ReLU6(inplace=True)
|
||||
)
|
||||
|
||||
|
||||
def conv_1x1_bn(inp, oup):
|
||||
return nn.Sequential(
|
||||
nn.Conv2d(inp, oup, 1, 1, 0, bias=False),
|
||||
nn.BatchNorm2d(oup),
|
||||
nn.ReLU6(inplace=True)
|
||||
)
|
||||
|
||||
|
||||
#------------------------------------------------------------------------------
|
||||
# Class of Inverted Residual block
|
||||
#------------------------------------------------------------------------------
|
||||
|
||||
class InvertedResidual(nn.Module):
|
||||
def __init__(self, inp, oup, stride, expansion, dilation=1):
|
||||
super(InvertedResidual, self).__init__()
|
||||
self.stride = stride
|
||||
assert stride in [1, 2]
|
||||
|
||||
hidden_dim = round(inp * expansion)
|
||||
self.use_res_connect = self.stride == 1 and inp == oup
|
||||
|
||||
if expansion == 1:
|
||||
self.conv = nn.Sequential(
|
||||
# dw
|
||||
nn.Conv2d(hidden_dim, hidden_dim, 3, stride, 1, groups=hidden_dim, dilation=dilation, bias=False),
|
||||
nn.BatchNorm2d(hidden_dim),
|
||||
nn.ReLU6(inplace=True),
|
||||
# pw-linear
|
||||
nn.Conv2d(hidden_dim, oup, 1, 1, 0, bias=False),
|
||||
nn.BatchNorm2d(oup),
|
||||
)
|
||||
else:
|
||||
self.conv = nn.Sequential(
|
||||
# pw
|
||||
nn.Conv2d(inp, hidden_dim, 1, 1, 0, bias=False),
|
||||
nn.BatchNorm2d(hidden_dim),
|
||||
nn.ReLU6(inplace=True),
|
||||
# dw
|
||||
nn.Conv2d(hidden_dim, hidden_dim, 3, stride, 1, groups=hidden_dim, dilation=dilation, bias=False),
|
||||
nn.BatchNorm2d(hidden_dim),
|
||||
nn.ReLU6(inplace=True),
|
||||
# pw-linear
|
||||
nn.Conv2d(hidden_dim, oup, 1, 1, 0, bias=False),
|
||||
nn.BatchNorm2d(oup),
|
||||
)
|
||||
|
||||
def forward(self, x):
|
||||
if self.use_res_connect:
|
||||
return x + self.conv(x)
|
||||
else:
|
||||
return self.conv(x)
|
||||
|
||||
|
||||
#------------------------------------------------------------------------------
|
||||
# Class of MobileNetV2
|
||||
#------------------------------------------------------------------------------
|
||||
|
||||
class MobileNetV2(nn.Module):
|
||||
def __init__(self, in_channels, alpha=1.0, expansion=6, num_classes=1000):
|
||||
super(MobileNetV2, self).__init__()
|
||||
self.in_channels = in_channels
|
||||
self.num_classes = num_classes
|
||||
input_channel = 32
|
||||
last_channel = 1280
|
||||
interverted_residual_setting = [
|
||||
# t, c, n, s
|
||||
[1 , 16, 1, 1],
|
||||
[expansion, 24, 2, 2],
|
||||
[expansion, 32, 3, 2],
|
||||
[expansion, 64, 4, 2],
|
||||
[expansion, 96, 3, 1],
|
||||
[expansion, 160, 3, 2],
|
||||
[expansion, 320, 1, 1],
|
||||
]
|
||||
|
||||
# building first layer
|
||||
input_channel = _make_divisible(input_channel*alpha, 8)
|
||||
self.last_channel = _make_divisible(last_channel*alpha, 8) if alpha > 1.0 else last_channel
|
||||
self.features = [conv_bn(self.in_channels, input_channel, 2)]
|
||||
|
||||
# building inverted residual blocks
|
||||
for t, c, n, s in interverted_residual_setting:
|
||||
output_channel = _make_divisible(int(c*alpha), 8)
|
||||
for i in range(n):
|
||||
if i == 0:
|
||||
self.features.append(InvertedResidual(input_channel, output_channel, s, expansion=t))
|
||||
else:
|
||||
self.features.append(InvertedResidual(input_channel, output_channel, 1, expansion=t))
|
||||
input_channel = output_channel
|
||||
|
||||
# building last several layers
|
||||
self.features.append(conv_1x1_bn(input_channel, self.last_channel))
|
||||
|
||||
# make it nn.Sequential
|
||||
self.features = nn.Sequential(*self.features)
|
||||
|
||||
# building classifier
|
||||
if self.num_classes is not None:
|
||||
self.classifier = nn.Sequential(
|
||||
nn.Dropout(0.2),
|
||||
nn.Linear(self.last_channel, num_classes),
|
||||
)
|
||||
|
||||
# Initialize weights
|
||||
self._init_weights()
|
||||
|
||||
def forward(self, x):
|
||||
# Stage1
|
||||
x = self.features[0](x)
|
||||
x = self.features[1](x)
|
||||
# Stage2
|
||||
x = self.features[2](x)
|
||||
x = self.features[3](x)
|
||||
# Stage3
|
||||
x = self.features[4](x)
|
||||
x = self.features[5](x)
|
||||
x = self.features[6](x)
|
||||
# Stage4
|
||||
x = self.features[7](x)
|
||||
x = self.features[8](x)
|
||||
x = self.features[9](x)
|
||||
x = self.features[10](x)
|
||||
x = self.features[11](x)
|
||||
x = self.features[12](x)
|
||||
x = self.features[13](x)
|
||||
# Stage5
|
||||
x = self.features[14](x)
|
||||
x = self.features[15](x)
|
||||
x = self.features[16](x)
|
||||
x = self.features[17](x)
|
||||
x = self.features[18](x)
|
||||
|
||||
# Classification
|
||||
if self.num_classes is not None:
|
||||
x = x.mean(dim=(2,3))
|
||||
x = self.classifier(x)
|
||||
|
||||
# Output
|
||||
return x
|
||||
|
||||
def _load_pretrained_model(self, pretrained_file):
|
||||
pretrain_dict = torch.load(pretrained_file, map_location='cpu')
|
||||
model_dict = {}
|
||||
state_dict = self.state_dict()
|
||||
print("[MobileNetV2] Loading pretrained model...")
|
||||
for k, v in pretrain_dict.items():
|
||||
if k in state_dict:
|
||||
model_dict[k] = v
|
||||
else:
|
||||
print(k, "is ignored")
|
||||
state_dict.update(model_dict)
|
||||
self.load_state_dict(state_dict)
|
||||
|
||||
def _init_weights(self):
|
||||
for m in self.modules():
|
||||
if isinstance(m, nn.Conv2d):
|
||||
n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
|
||||
m.weight.data.normal_(0, math.sqrt(2. / n))
|
||||
if m.bias is not None:
|
||||
m.bias.data.zero_()
|
||||
elif isinstance(m, nn.BatchNorm2d):
|
||||
m.weight.data.fill_(1)
|
||||
m.bias.data.zero_()
|
||||
elif isinstance(m, nn.Linear):
|
||||
n = m.weight.size(1)
|
||||
m.weight.data.normal_(0, 0.01)
|
||||
m.bias.data.zero_()
|
||||
@@ -0,0 +1,82 @@
|
||||
import os
|
||||
from functools import reduce
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
from .mobilenetv2 import MobileNetV2
|
||||
|
||||
|
||||
class BaseBackbone(nn.Module):
|
||||
""" Superclass of Replaceable Backbone Model for Semantic Estimation
|
||||
"""
|
||||
|
||||
def __init__(self, in_channels):
|
||||
super(BaseBackbone, self).__init__()
|
||||
self.in_channels = in_channels
|
||||
|
||||
self.model = None
|
||||
self.enc_channels = []
|
||||
|
||||
def forward(self, x):
|
||||
raise NotImplementedError
|
||||
|
||||
def load_pretrained_ckpt(self):
|
||||
raise NotImplementedError
|
||||
|
||||
|
||||
class MobileNetV2Backbone(BaseBackbone):
|
||||
""" MobileNetV2 Backbone
|
||||
"""
|
||||
|
||||
def __init__(self, in_channels):
|
||||
super(MobileNetV2Backbone, self).__init__(in_channels)
|
||||
|
||||
self.model = MobileNetV2(self.in_channels, alpha=1.0, expansion=6, num_classes=None)
|
||||
self.enc_channels = [16, 24, 32, 96, 1280]
|
||||
|
||||
def forward(self, x):
|
||||
# x = reduce(lambda x, n: self.model.features[n](x), list(range(0, 2)), x)
|
||||
x = self.model.features[0](x)
|
||||
x = self.model.features[1](x)
|
||||
enc2x = x
|
||||
|
||||
# x = reduce(lambda x, n: self.model.features[n](x), list(range(2, 4)), x)
|
||||
x = self.model.features[2](x)
|
||||
x = self.model.features[3](x)
|
||||
enc4x = x
|
||||
|
||||
# x = reduce(lambda x, n: self.model.features[n](x), list(range(4, 7)), x)
|
||||
x = self.model.features[4](x)
|
||||
x = self.model.features[5](x)
|
||||
x = self.model.features[6](x)
|
||||
enc8x = x
|
||||
|
||||
# x = reduce(lambda x, n: self.model.features[n](x), list(range(7, 14)), x)
|
||||
x = self.model.features[7](x)
|
||||
x = self.model.features[8](x)
|
||||
x = self.model.features[9](x)
|
||||
x = self.model.features[10](x)
|
||||
x = self.model.features[11](x)
|
||||
x = self.model.features[12](x)
|
||||
x = self.model.features[13](x)
|
||||
enc16x = x
|
||||
|
||||
# x = reduce(lambda x, n: self.model.features[n](x), list(range(14, 19)), x)
|
||||
x = self.model.features[14](x)
|
||||
x = self.model.features[15](x)
|
||||
x = self.model.features[16](x)
|
||||
x = self.model.features[17](x)
|
||||
x = self.model.features[18](x)
|
||||
enc32x = x
|
||||
return [enc2x, enc4x, enc8x, enc16x, enc32x]
|
||||
|
||||
def load_pretrained_ckpt(self):
|
||||
# the pre-trained model is provided by https://github.com/thuyngch/Human-Segmentation-PyTorch
|
||||
ckpt_path = './pretrained/mobilenetv2_human_seg.ckpt'
|
||||
if not os.path.exists(ckpt_path):
|
||||
print('cannot find the pretrained mobilenetv2 backbone')
|
||||
exit()
|
||||
|
||||
ckpt = torch.load(ckpt_path)
|
||||
self.model.load_state_dict(ckpt)
|
||||
@@ -0,0 +1,255 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
from .backbones import SUPPORTED_BACKBONES
|
||||
|
||||
|
||||
#------------------------------------------------------------------------------
|
||||
# MODNet Basic Modules
|
||||
#------------------------------------------------------------------------------
|
||||
|
||||
class IBNorm(nn.Module):
|
||||
""" Combine Instance Norm and Batch Norm into One Layer
|
||||
"""
|
||||
|
||||
def __init__(self, in_channels):
|
||||
super(IBNorm, self).__init__()
|
||||
in_channels = in_channels
|
||||
self.bnorm_channels = int(in_channels / 2)
|
||||
self.inorm_channels = in_channels - self.bnorm_channels
|
||||
|
||||
self.bnorm = nn.BatchNorm2d(self.bnorm_channels, affine=True)
|
||||
self.inorm = nn.InstanceNorm2d(self.inorm_channels, affine=False)
|
||||
|
||||
def forward(self, x):
|
||||
bn_x = self.bnorm(x[:, :self.bnorm_channels, ...].contiguous())
|
||||
in_x = self.inorm(x[:, self.bnorm_channels:, ...].contiguous())
|
||||
|
||||
return torch.cat((bn_x, in_x), 1)
|
||||
|
||||
|
||||
class Conv2dIBNormRelu(nn.Module):
|
||||
""" Convolution + IBNorm + ReLu
|
||||
"""
|
||||
|
||||
def __init__(self, in_channels, out_channels, kernel_size,
|
||||
stride=1, padding=0, dilation=1, groups=1, bias=True,
|
||||
with_ibn=True, with_relu=True):
|
||||
super(Conv2dIBNormRelu, self).__init__()
|
||||
|
||||
layers = [
|
||||
nn.Conv2d(in_channels, out_channels, kernel_size,
|
||||
stride=stride, padding=padding, dilation=dilation,
|
||||
groups=groups, bias=bias)
|
||||
]
|
||||
|
||||
if with_ibn:
|
||||
layers.append(IBNorm(out_channels))
|
||||
if with_relu:
|
||||
layers.append(nn.ReLU(inplace=True))
|
||||
|
||||
self.layers = nn.Sequential(*layers)
|
||||
|
||||
def forward(self, x):
|
||||
return self.layers(x)
|
||||
|
||||
|
||||
class SEBlock(nn.Module):
|
||||
""" SE Block Proposed in https://arxiv.org/pdf/1709.01507.pdf
|
||||
"""
|
||||
|
||||
def __init__(self, in_channels, out_channels, reduction=1):
|
||||
super(SEBlock, self).__init__()
|
||||
self.pool = nn.AdaptiveAvgPool2d(1)
|
||||
self.fc = nn.Sequential(
|
||||
nn.Linear(in_channels, int(in_channels // reduction), bias=False),
|
||||
nn.ReLU(inplace=True),
|
||||
nn.Linear(int(in_channels // reduction), out_channels, bias=False),
|
||||
nn.Sigmoid()
|
||||
)
|
||||
|
||||
def forward(self, x):
|
||||
b, c, _, _ = x.size()
|
||||
w = self.pool(x).view(b, c)
|
||||
w = self.fc(w).view(b, c, 1, 1)
|
||||
|
||||
return x * w.expand_as(x)
|
||||
|
||||
|
||||
#------------------------------------------------------------------------------
|
||||
# MODNet Branches
|
||||
#------------------------------------------------------------------------------
|
||||
|
||||
class LRBranch(nn.Module):
|
||||
""" Low Resolution Branch of MODNet
|
||||
"""
|
||||
|
||||
def __init__(self, backbone):
|
||||
super(LRBranch, self).__init__()
|
||||
|
||||
enc_channels = backbone.enc_channels
|
||||
|
||||
self.backbone = backbone
|
||||
self.se_block = SEBlock(enc_channels[4], enc_channels[4], reduction=4)
|
||||
self.conv_lr16x = Conv2dIBNormRelu(enc_channels[4], enc_channels[3], 5, stride=1, padding=2)
|
||||
self.conv_lr8x = Conv2dIBNormRelu(enc_channels[3], enc_channels[2], 5, stride=1, padding=2)
|
||||
self.conv_lr = Conv2dIBNormRelu(enc_channels[2], 1, kernel_size=3, stride=2, padding=1, with_ibn=False, with_relu=False)
|
||||
|
||||
def forward(self, img, inference):
|
||||
enc_features = self.backbone.forward(img)
|
||||
enc2x, enc4x, enc32x = enc_features[0], enc_features[1], enc_features[4]
|
||||
|
||||
enc32x = self.se_block(enc32x)
|
||||
lr16x = F.interpolate(enc32x, scale_factor=2, mode='bilinear', align_corners=False)
|
||||
lr16x = self.conv_lr16x(lr16x)
|
||||
lr8x = F.interpolate(lr16x, scale_factor=2, mode='bilinear', align_corners=False)
|
||||
lr8x = self.conv_lr8x(lr8x)
|
||||
|
||||
pred_semantic = None
|
||||
if not inference:
|
||||
lr = self.conv_lr(lr8x)
|
||||
pred_semantic = torch.sigmoid(lr)
|
||||
|
||||
return pred_semantic, lr8x, [enc2x, enc4x]
|
||||
|
||||
|
||||
class HRBranch(nn.Module):
|
||||
""" High Resolution Branch of MODNet
|
||||
"""
|
||||
|
||||
def __init__(self, hr_channels, enc_channels):
|
||||
super(HRBranch, self).__init__()
|
||||
|
||||
self.tohr_enc2x = Conv2dIBNormRelu(enc_channels[0], hr_channels, 1, stride=1, padding=0)
|
||||
self.conv_enc2x = Conv2dIBNormRelu(hr_channels + 3, hr_channels, 3, stride=2, padding=1)
|
||||
|
||||
self.tohr_enc4x = Conv2dIBNormRelu(enc_channels[1], hr_channels, 1, stride=1, padding=0)
|
||||
self.conv_enc4x = Conv2dIBNormRelu(2 * hr_channels, 2 * hr_channels, 3, stride=1, padding=1)
|
||||
|
||||
self.conv_hr4x = nn.Sequential(
|
||||
Conv2dIBNormRelu(3 * hr_channels + 3, 2 * hr_channels, 3, stride=1, padding=1),
|
||||
Conv2dIBNormRelu(2 * hr_channels, 2 * hr_channels, 3, stride=1, padding=1),
|
||||
Conv2dIBNormRelu(2 * hr_channels, hr_channels, 3, stride=1, padding=1),
|
||||
)
|
||||
|
||||
self.conv_hr2x = nn.Sequential(
|
||||
Conv2dIBNormRelu(2 * hr_channels, 2 * hr_channels, 3, stride=1, padding=1),
|
||||
Conv2dIBNormRelu(2 * hr_channels, hr_channels, 3, stride=1, padding=1),
|
||||
Conv2dIBNormRelu(hr_channels, hr_channels, 3, stride=1, padding=1),
|
||||
Conv2dIBNormRelu(hr_channels, hr_channels, 3, stride=1, padding=1),
|
||||
)
|
||||
|
||||
self.conv_hr = nn.Sequential(
|
||||
Conv2dIBNormRelu(hr_channels + 3, hr_channels, 3, stride=1, padding=1),
|
||||
Conv2dIBNormRelu(hr_channels, 1, kernel_size=1, stride=1, padding=0, with_ibn=False, with_relu=False),
|
||||
)
|
||||
|
||||
def forward(self, img, enc2x, enc4x, lr8x, inference):
|
||||
img2x = F.interpolate(img, scale_factor=1/2, mode='bilinear', align_corners=False)
|
||||
img4x = F.interpolate(img, scale_factor=1/4, mode='bilinear', align_corners=False)
|
||||
|
||||
enc2x = self.tohr_enc2x(enc2x)
|
||||
hr4x = self.conv_enc2x(torch.cat((img2x, enc2x), dim=1))
|
||||
|
||||
enc4x = self.tohr_enc4x(enc4x)
|
||||
hr4x = self.conv_enc4x(torch.cat((hr4x, enc4x), dim=1))
|
||||
|
||||
lr4x = F.interpolate(lr8x, scale_factor=2, mode='bilinear', align_corners=False)
|
||||
hr4x = self.conv_hr4x(torch.cat((hr4x, lr4x, img4x), dim=1))
|
||||
|
||||
hr2x = F.interpolate(hr4x, scale_factor=2, mode='bilinear', align_corners=False)
|
||||
hr2x = self.conv_hr2x(torch.cat((hr2x, enc2x), dim=1))
|
||||
|
||||
pred_detail = None
|
||||
if not inference:
|
||||
hr = F.interpolate(hr2x, scale_factor=2, mode='bilinear', align_corners=False)
|
||||
hr = self.conv_hr(torch.cat((hr, img), dim=1))
|
||||
pred_detail = torch.sigmoid(hr)
|
||||
|
||||
return pred_detail, hr2x
|
||||
|
||||
|
||||
class FusionBranch(nn.Module):
|
||||
""" Fusion Branch of MODNet
|
||||
"""
|
||||
|
||||
def __init__(self, hr_channels, enc_channels):
|
||||
super(FusionBranch, self).__init__()
|
||||
self.conv_lr4x = Conv2dIBNormRelu(enc_channels[2], hr_channels, 5, stride=1, padding=2)
|
||||
|
||||
self.conv_f2x = Conv2dIBNormRelu(2 * hr_channels, hr_channels, 3, stride=1, padding=1)
|
||||
self.conv_f = nn.Sequential(
|
||||
Conv2dIBNormRelu(hr_channels + 3, int(hr_channels / 2), 3, stride=1, padding=1),
|
||||
Conv2dIBNormRelu(int(hr_channels / 2), 1, 1, stride=1, padding=0, with_ibn=False, with_relu=False),
|
||||
)
|
||||
|
||||
def forward(self, img, lr8x, hr2x):
|
||||
lr4x = F.interpolate(lr8x, scale_factor=2, mode='bilinear', align_corners=False)
|
||||
lr4x = self.conv_lr4x(lr4x)
|
||||
lr2x = F.interpolate(lr4x, scale_factor=2, mode='bilinear', align_corners=False)
|
||||
|
||||
f2x = self.conv_f2x(torch.cat((lr2x, hr2x), dim=1))
|
||||
f = F.interpolate(f2x, scale_factor=2, mode='bilinear', align_corners=False)
|
||||
f = self.conv_f(torch.cat((f, img), dim=1))
|
||||
pred_matte = torch.sigmoid(f)
|
||||
|
||||
return pred_matte
|
||||
|
||||
|
||||
#------------------------------------------------------------------------------
|
||||
# MODNet
|
||||
#------------------------------------------------------------------------------
|
||||
|
||||
class MODNet(nn.Module):
|
||||
""" Architecture of MODNet
|
||||
"""
|
||||
|
||||
def __init__(self, in_channels=3, hr_channels=32, backbone_arch='mobilenetv2', backbone_pretrained=True):
|
||||
super(MODNet, self).__init__()
|
||||
|
||||
self.in_channels = in_channels
|
||||
self.hr_channels = hr_channels
|
||||
self.backbone_arch = backbone_arch
|
||||
self.backbone_pretrained = backbone_pretrained
|
||||
|
||||
self.backbone = SUPPORTED_BACKBONES[self.backbone_arch](self.in_channels)
|
||||
|
||||
self.lr_branch = LRBranch(self.backbone)
|
||||
self.hr_branch = HRBranch(self.hr_channels, self.backbone.enc_channels)
|
||||
self.f_branch = FusionBranch(self.hr_channels, self.backbone.enc_channels)
|
||||
|
||||
for m in self.modules():
|
||||
if isinstance(m, nn.Conv2d):
|
||||
self._init_conv(m)
|
||||
elif isinstance(m, nn.BatchNorm2d) or isinstance(m, nn.InstanceNorm2d):
|
||||
self._init_norm(m)
|
||||
|
||||
if self.backbone_pretrained:
|
||||
self.backbone.load_pretrained_ckpt()
|
||||
|
||||
def forward(self, img, inference):
|
||||
pred_semantic, lr8x, [enc2x, enc4x] = self.lr_branch(img, inference)
|
||||
pred_detail, hr2x = self.hr_branch(img, enc2x, enc4x, lr8x, inference)
|
||||
pred_matte = self.f_branch(img, lr8x, hr2x)
|
||||
|
||||
return pred_semantic, pred_detail, pred_matte
|
||||
|
||||
def freeze_norm(self):
|
||||
norm_types = [nn.BatchNorm2d, nn.InstanceNorm2d]
|
||||
for m in self.modules():
|
||||
for n in norm_types:
|
||||
if isinstance(m, n):
|
||||
m.eval()
|
||||
continue
|
||||
|
||||
def _init_conv(self, conv):
|
||||
nn.init.kaiming_uniform_(
|
||||
conv.weight, a=0, mode='fan_in', nonlinearity='relu')
|
||||
if conv.bias is not None:
|
||||
nn.init.constant_(conv.bias, 0)
|
||||
|
||||
def _init_norm(self, norm):
|
||||
if norm.weight is not None:
|
||||
nn.init.constant_(norm.weight, 1)
|
||||
nn.init.constant_(norm.bias, 0)
|
||||
@@ -0,0 +1,299 @@
|
||||
import math
|
||||
import scipy
|
||||
import numpy as np
|
||||
from scipy.ndimage import grey_dilation, grey_erosion
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
|
||||
__all__ = [
|
||||
'supervised_training_iter',
|
||||
'soc_adaptation_iter',
|
||||
]
|
||||
|
||||
|
||||
# ----------------------------------------------------------------------------------
|
||||
# Tool Classes/Functions
|
||||
# ----------------------------------------------------------------------------------
|
||||
|
||||
class GaussianBlurLayer(nn.Module):
|
||||
""" Add Gaussian Blur to a 4D tensors
|
||||
This layer takes a 4D tensor of {N, C, H, W} as input.
|
||||
The Gaussian blur will be performed in given channel number (C) splitly.
|
||||
"""
|
||||
|
||||
def __init__(self, channels, kernel_size):
|
||||
"""
|
||||
Arguments:
|
||||
channels (int): Channel for input tensor
|
||||
kernel_size (int): Size of the kernel used in blurring
|
||||
"""
|
||||
|
||||
super(GaussianBlurLayer, self).__init__()
|
||||
self.channels = channels
|
||||
self.kernel_size = kernel_size
|
||||
assert self.kernel_size % 2 != 0
|
||||
|
||||
self.op = nn.Sequential(
|
||||
nn.ReflectionPad2d(math.floor(self.kernel_size / 2)),
|
||||
nn.Conv2d(channels, channels, self.kernel_size,
|
||||
stride=1, padding=0, bias=None, groups=channels)
|
||||
)
|
||||
|
||||
self._init_kernel()
|
||||
|
||||
def forward(self, x):
|
||||
"""
|
||||
Arguments:
|
||||
x (torch.Tensor): input 4D tensor
|
||||
Returns:
|
||||
torch.Tensor: Blurred version of the input
|
||||
"""
|
||||
|
||||
if not len(list(x.shape)) == 4:
|
||||
print('\'GaussianBlurLayer\' requires a 4D tensor as input\n')
|
||||
exit()
|
||||
elif not x.shape[1] == self.channels:
|
||||
print('In \'GaussianBlurLayer\', the required channel ({0}) is'
|
||||
'not the same as input ({1})\n'.format(self.channels, x.shape[1]))
|
||||
exit()
|
||||
|
||||
return self.op(x)
|
||||
|
||||
def _init_kernel(self):
|
||||
sigma = 0.3 * ((self.kernel_size - 1) * 0.5 - 1) + 0.8
|
||||
|
||||
n = np.zeros((self.kernel_size, self.kernel_size))
|
||||
i = math.floor(self.kernel_size / 2)
|
||||
n[i, i] = 1
|
||||
kernel = scipy.ndimage.gaussian_filter(n, sigma)
|
||||
|
||||
for name, param in self.named_parameters():
|
||||
param.data.copy_(torch.from_numpy(kernel))
|
||||
|
||||
# ----------------------------------------------------------------------------------
|
||||
|
||||
|
||||
# ----------------------------------------------------------------------------------
|
||||
# MODNet Training Functions
|
||||
# ----------------------------------------------------------------------------------
|
||||
|
||||
blurer = GaussianBlurLayer(1, 3).cuda()
|
||||
|
||||
|
||||
def supervised_training_iter(
|
||||
modnet, optimizer, image, trimap, gt_matte,
|
||||
semantic_scale=10.0, detail_scale=10.0, matte_scale=1.0):
|
||||
""" Supervised training iteration of MODNet
|
||||
This function trains MODNet for one iteration in a labeled dataset.
|
||||
|
||||
Arguments:
|
||||
modnet (torch.nn.Module): instance of MODNet
|
||||
optimizer (torch.optim.Optimizer): optimizer for supervised training
|
||||
image (torch.autograd.Variable): input RGB image
|
||||
its pixel values should be normalized
|
||||
trimap (torch.autograd.Variable): trimap used to calculate the losses
|
||||
its pixel values can be 0, 0.5, or 1
|
||||
(foreground=1, background=0, unknown=0.5)
|
||||
gt_matte (torch.autograd.Variable): ground truth alpha matte
|
||||
its pixel values are between [0, 1]
|
||||
semantic_scale (float): scale of the semantic loss
|
||||
NOTE: please adjust according to your dataset
|
||||
detail_scale (float): scale of the detail loss
|
||||
NOTE: please adjust according to your dataset
|
||||
matte_scale (float): scale of the matte loss
|
||||
NOTE: please adjust according to your dataset
|
||||
|
||||
Returns:
|
||||
semantic_loss (torch.Tensor): loss of the semantic estimation [Low-Resolution (LR) Branch]
|
||||
detail_loss (torch.Tensor): loss of the detail prediction [High-Resolution (HR) Branch]
|
||||
matte_loss (torch.Tensor): loss of the semantic-detail fusion [Fusion Branch]
|
||||
|
||||
Example:
|
||||
import torch
|
||||
from src.models.modnet import MODNet
|
||||
from src.trainer import supervised_training_iter
|
||||
|
||||
bs = 16 # batch size
|
||||
lr = 0.01 # learn rate
|
||||
epochs = 40 # total epochs
|
||||
|
||||
modnet = torch.nn.DataParallel(MODNet()).cuda()
|
||||
optimizer = torch.optim.SGD(modnet.parameters(), lr=lr, momentum=0.9)
|
||||
lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=int(0.25 * epochs), gamma=0.1)
|
||||
|
||||
dataloader = CREATE_YOUR_DATALOADER(bs) # NOTE: please finish this function
|
||||
|
||||
for epoch in range(0, epochs):
|
||||
for idx, (image, trimap, gt_matte) in enumerate(dataloader):
|
||||
semantic_loss, detail_loss, matte_loss = \
|
||||
supervised_training_iter(modnet, optimizer, image, trimap, gt_matte)
|
||||
lr_scheduler.step()
|
||||
"""
|
||||
|
||||
global blurer
|
||||
|
||||
# set the model to train mode and clear the optimizer
|
||||
modnet.train()
|
||||
optimizer.zero_grad()
|
||||
|
||||
# forward the model
|
||||
pred_semantic, pred_detail, pred_matte = modnet(image, False)
|
||||
|
||||
# calculate the boundary mask from the trimap
|
||||
boundaries = (trimap < 0.5) + (trimap > 0.5)
|
||||
|
||||
# calculate the semantic loss
|
||||
gt_semantic = F.interpolate(gt_matte, scale_factor=1/16, mode='bilinear')
|
||||
gt_semantic = blurer(gt_semantic)
|
||||
semantic_loss = torch.mean(F.mse_loss(pred_semantic, gt_semantic))
|
||||
semantic_loss = semantic_scale * semantic_loss
|
||||
|
||||
# calculate the detail loss
|
||||
pred_boundary_detail = torch.where(boundaries, trimap, pred_detail)
|
||||
gt_detail = torch.where(boundaries, trimap, gt_matte)
|
||||
detail_loss = torch.mean(F.l1_loss(pred_boundary_detail, gt_detail))
|
||||
detail_loss = detail_scale * detail_loss
|
||||
|
||||
# calculate the matte loss
|
||||
pred_boundary_matte = torch.where(boundaries, trimap, pred_matte)
|
||||
matte_l1_loss = F.l1_loss(pred_matte, gt_matte) + 4.0 * F.l1_loss(pred_boundary_matte, gt_matte)
|
||||
matte_compositional_loss = F.l1_loss(image * pred_matte, image * gt_matte) \
|
||||
+ 4.0 * F.l1_loss(image * pred_boundary_matte, image * gt_matte)
|
||||
matte_loss = torch.mean(matte_l1_loss + matte_compositional_loss)
|
||||
matte_loss = matte_scale * matte_loss
|
||||
|
||||
# calculate the final loss, backward the loss, and update the model
|
||||
loss = semantic_loss + detail_loss + matte_loss
|
||||
loss.backward()
|
||||
optimizer.step()
|
||||
|
||||
# for test
|
||||
return semantic_loss, detail_loss, matte_loss
|
||||
|
||||
|
||||
def soc_adaptation_iter(
|
||||
modnet, backup_modnet, optimizer, image,
|
||||
soc_semantic_scale=100.0, soc_detail_scale=1.0):
|
||||
""" Self-Supervised sub-objective consistency (SOC) adaptation iteration of MODNet
|
||||
This function fine-tunes MODNet for one iteration in an unlabeled dataset.
|
||||
Note that SOC can only fine-tune a converged MODNet, i.e., MODNet that has been
|
||||
trained in a labeled dataset.
|
||||
|
||||
Arguments:
|
||||
modnet (torch.nn.Module): instance of MODNet
|
||||
backup_modnet (torch.nn.Module): backup of the trained MODNet
|
||||
optimizer (torch.optim.Optimizer): optimizer for self-supervised SOC
|
||||
image (torch.autograd.Variable): input RGB image
|
||||
its pixel values should be normalized
|
||||
soc_semantic_scale (float): scale of the SOC semantic loss
|
||||
NOTE: please adjust according to your dataset
|
||||
soc_detail_scale (float): scale of the SOC detail loss
|
||||
NOTE: please adjust according to your dataset
|
||||
|
||||
Returns:
|
||||
soc_semantic_loss (torch.Tensor): loss of the semantic SOC
|
||||
soc_detail_loss (torch.Tensor): loss of the detail SOC
|
||||
|
||||
Example:
|
||||
import copy
|
||||
import torch
|
||||
from src.models.modnet import MODNet
|
||||
from src.trainer import soc_adaptation_iter
|
||||
|
||||
bs = 1 # batch size
|
||||
lr = 0.00001 # learn rate
|
||||
epochs = 10 # total epochs
|
||||
|
||||
modnet = torch.nn.DataParallel(MODNet()).cuda()
|
||||
modnet = LOAD_TRAINED_CKPT() # NOTE: please finish this function
|
||||
|
||||
optimizer = torch.optim.Adam(modnet.parameters(), lr=lr, betas=(0.9, 0.99))
|
||||
dataloader = CREATE_YOUR_DATALOADER(bs) # NOTE: please finish this function
|
||||
|
||||
for epoch in range(0, epochs):
|
||||
backup_modnet = copy.deepcopy(modnet)
|
||||
for idx, (image) in enumerate(dataloader):
|
||||
soc_semantic_loss, soc_detail_loss = \
|
||||
soc_adaptation_iter(modnet, backup_modnet, optimizer, image)
|
||||
"""
|
||||
|
||||
global blurer
|
||||
|
||||
# set the backup model to eval mode
|
||||
backup_modnet.eval()
|
||||
|
||||
# set the main model to train mode and freeze its norm layers
|
||||
modnet.train()
|
||||
modnet.module.freeze_norm()
|
||||
|
||||
# clear the optimizer
|
||||
optimizer.zero_grad()
|
||||
|
||||
# forward the main model
|
||||
pred_semantic, pred_detail, pred_matte = modnet(image, False)
|
||||
|
||||
# forward the backup model
|
||||
with torch.no_grad():
|
||||
_, pred_backup_detail, pred_backup_matte = backup_modnet(image, False)
|
||||
|
||||
# calculate the boundary mask from `pred_matte` and `pred_semantic`
|
||||
pred_matte_fg = (pred_matte.detach() > 0.1).float()
|
||||
pred_semantic_fg = (pred_semantic.detach() > 0.1).float()
|
||||
pred_semantic_fg = F.interpolate(pred_semantic_fg, scale_factor=16, mode='bilinear')
|
||||
pred_fg = pred_matte_fg * pred_semantic_fg
|
||||
|
||||
n, c, h, w = pred_matte.shape
|
||||
np_pred_fg = pred_fg.data.cpu().numpy()
|
||||
np_boundaries = np.zeros([n, c, h, w])
|
||||
for sdx in range(0, n):
|
||||
sample_np_boundaries = np_boundaries[sdx, 0, ...]
|
||||
sample_np_pred_fg = np_pred_fg[sdx, 0, ...]
|
||||
|
||||
side = int((h + w) / 2 * 0.05)
|
||||
dilated = grey_dilation(sample_np_pred_fg, size=(side, side))
|
||||
eroded = grey_erosion(sample_np_pred_fg, size=(side, side))
|
||||
|
||||
sample_np_boundaries[np.where(dilated - eroded != 0)] = 1
|
||||
np_boundaries[sdx, 0, ...] = sample_np_boundaries
|
||||
|
||||
boundaries = torch.tensor(np_boundaries).float().cuda()
|
||||
|
||||
# sub-objectives consistency between `pred_semantic` and `pred_matte`
|
||||
# generate pseudo ground truth for `pred_semantic`
|
||||
downsampled_pred_matte = blurer(F.interpolate(pred_matte, scale_factor=1/16, mode='bilinear'))
|
||||
pseudo_gt_semantic = downsampled_pred_matte.detach()
|
||||
pseudo_gt_semantic = pseudo_gt_semantic * (pseudo_gt_semantic > 0.01).float()
|
||||
|
||||
# generate pseudo ground truth for `pred_matte`
|
||||
pseudo_gt_matte = pred_semantic.detach()
|
||||
pseudo_gt_matte = pseudo_gt_matte * (pseudo_gt_matte > 0.01).float()
|
||||
|
||||
# calculate the SOC semantic loss
|
||||
soc_semantic_loss = F.mse_loss(pred_semantic, pseudo_gt_semantic) + F.mse_loss(downsampled_pred_matte, pseudo_gt_matte)
|
||||
soc_semantic_loss = soc_semantic_scale * torch.mean(soc_semantic_loss)
|
||||
|
||||
# NOTE: using the formulas in our paper to calculate the following losses has similar results
|
||||
# sub-objectives consistency between `pred_detail` and `pred_backup_detail` (on boundaries only)
|
||||
backup_detail_loss = boundaries * F.l1_loss(pred_detail, pred_backup_detail, reduction='none')
|
||||
backup_detail_loss = torch.sum(backup_detail_loss, dim=(1,2,3)) / torch.sum(boundaries, dim=(1,2,3))
|
||||
backup_detail_loss = torch.mean(backup_detail_loss)
|
||||
|
||||
# sub-objectives consistency between pred_matte` and `pred_backup_matte` (on boundaries only)
|
||||
backup_matte_loss = boundaries * F.l1_loss(pred_matte, pred_backup_matte, reduction='none')
|
||||
backup_matte_loss = torch.sum(backup_matte_loss, dim=(1,2,3)) / torch.sum(boundaries, dim=(1,2,3))
|
||||
backup_matte_loss = torch.mean(backup_matte_loss)
|
||||
|
||||
soc_detail_loss = soc_detail_scale * (backup_detail_loss + backup_matte_loss)
|
||||
|
||||
# calculate the final loss, backward the loss, and update the model
|
||||
loss = soc_semantic_loss + soc_detail_loss
|
||||
|
||||
loss.backward()
|
||||
optimizer.step()
|
||||
|
||||
return soc_semantic_loss, soc_detail_loss
|
||||
|
||||
# ----------------------------------------------------------------------------------
|
||||
@@ -0,0 +1,18 @@
|
||||
## MODNet - TorchScript Model
|
||||
|
||||
This TorchScript version of MODNet is provided by [@yarkable](https://github.com/yarkable) from the community.
|
||||
Please note that the PyTorch version required for this TorchScript export function is higher than the official MODNet code (torch>=1.2.0).
|
||||
|
||||
You can also download the TorchScript version of the official **Image Matting Model** from [this link](https://pan.baidu.com/s/1kOmmmbG7lSZiSmDdE7CaRw) with the exextraction code `dm9e`.
|
||||
|
||||
To export the TorchScript version of MODNet (assuming you are currently in project root directory):
|
||||
1. Download the pre-trained **Image Matting Model** from this [link](https://drive.google.com/drive/folders/1umYmlCulvIFNaqPjwod1SayFmSRHziyR?usp=sharing) and put the model into the folder `MODNet/pretrained/`.
|
||||
|
||||
2. Ensure your PyTorch version >= 1.2.0.
|
||||
|
||||
3. Export the TorchScript version of MODNet by:
|
||||
```shell
|
||||
python -m torchscript.export_torchscript \
|
||||
--ckpt-path=pretrained/modnet_photographic_portrait_matting.ckpt \
|
||||
--output-path=pretrained/modnet_photographic_portrait_matting.torchscript
|
||||
```
|
||||
@@ -0,0 +1,46 @@
|
||||
"""
|
||||
Export TorchScript model of MODNet
|
||||
|
||||
Arguments:
|
||||
--ckpt-path: path of the checkpoint that will be converted
|
||||
--output-path: path for saving the TorchScript model
|
||||
|
||||
Example:
|
||||
python export_torchscript.py \
|
||||
--ckpt-path=modnet_photographic_portrait_matting.ckpt \
|
||||
--output-path=modnet_photographic_portrait_matting.torchscript
|
||||
"""
|
||||
|
||||
import os
|
||||
import argparse
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
from . import modnet_torchscript
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
# define cmd arguments
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument('--ckpt-path', type=str, required=True, help='path of the checkpoint that will be converted')
|
||||
parser.add_argument('--output-path', type=str, required=True, help='path for saving the TorchScript model')
|
||||
args = parser.parse_args()
|
||||
|
||||
# check input arguments
|
||||
if not os.path.exists(args.ckpt_path):
|
||||
print(args.ckpt_path)
|
||||
print('Cannot find checkpoint path: {0}'.format(args.ckpt_path))
|
||||
exit()
|
||||
|
||||
# create MODNet and load the pre-trained ckpt
|
||||
modnet = modnet_torchscript.MODNet(backbone_pretrained=False)
|
||||
modnet = nn.DataParallel(modnet).cuda()
|
||||
state_dict = torch.load(args.ckpt_path)
|
||||
modnet.load_state_dict(state_dict)
|
||||
modnet.eval()
|
||||
|
||||
# export to TorchScript model
|
||||
scripted_model = torch.jit.script(modnet.module)
|
||||
torch.jit.save(scripted_model, os.path.join(args.output_path))
|
||||
@@ -0,0 +1,258 @@
|
||||
"""
|
||||
This file contains a modified version of the original file `modnet.py` without
|
||||
`pred_semantic` and `pred_details` as these both returns None when `inference=True`
|
||||
|
||||
And it does not contain `inference` argument which will make it easier to
|
||||
convert checkpoint to TorchScript model.
|
||||
"""
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
from src.models.backbones import SUPPORTED_BACKBONES
|
||||
|
||||
|
||||
#------------------------------------------------------------------------------
|
||||
# MODNet Basic Modules
|
||||
#------------------------------------------------------------------------------
|
||||
|
||||
class IBNorm(nn.Module):
|
||||
""" Combine Instance Norm and Batch Norm into One Layer
|
||||
"""
|
||||
|
||||
def __init__(self, in_channels):
|
||||
super(IBNorm, self).__init__()
|
||||
in_channels = in_channels
|
||||
self.bnorm_channels = int(in_channels / 2)
|
||||
self.inorm_channels = in_channels - self.bnorm_channels
|
||||
|
||||
self.bnorm = nn.BatchNorm2d(self.bnorm_channels, affine=True)
|
||||
self.inorm = nn.InstanceNorm2d(self.inorm_channels, affine=False)
|
||||
|
||||
def forward(self, x):
|
||||
bn_x = self.bnorm(x[:, :self.bnorm_channels, ...].contiguous())
|
||||
in_x = self.inorm(x[:, self.bnorm_channels:, ...].contiguous())
|
||||
|
||||
return torch.cat((bn_x, in_x), 1)
|
||||
|
||||
|
||||
class Conv2dIBNormRelu(nn.Module):
|
||||
""" Convolution + IBNorm + ReLu
|
||||
"""
|
||||
|
||||
def __init__(self, in_channels, out_channels, kernel_size,
|
||||
stride=1, padding=0, dilation=1, groups=1, bias=True,
|
||||
with_ibn=True, with_relu=True):
|
||||
super(Conv2dIBNormRelu, self).__init__()
|
||||
|
||||
layers = [
|
||||
nn.Conv2d(in_channels, out_channels, kernel_size,
|
||||
stride=stride, padding=padding, dilation=dilation,
|
||||
groups=groups, bias=bias)
|
||||
]
|
||||
|
||||
if with_ibn:
|
||||
layers.append(IBNorm(out_channels))
|
||||
if with_relu:
|
||||
layers.append(nn.ReLU(inplace=True))
|
||||
|
||||
self.layers = nn.Sequential(*layers)
|
||||
|
||||
def forward(self, x):
|
||||
return self.layers(x)
|
||||
|
||||
|
||||
class SEBlock(nn.Module):
|
||||
""" SE Block Proposed in https://arxiv.org/pdf/1709.01507.pdf
|
||||
"""
|
||||
|
||||
def __init__(self, in_channels, out_channels, reduction=1):
|
||||
super(SEBlock, self).__init__()
|
||||
self.pool = nn.AdaptiveAvgPool2d(1)
|
||||
self.fc = nn.Sequential(
|
||||
nn.Linear(in_channels, int(in_channels // reduction), bias=False),
|
||||
nn.ReLU(inplace=True),
|
||||
nn.Linear(int(in_channels // reduction), out_channels, bias=False),
|
||||
nn.Sigmoid()
|
||||
)
|
||||
|
||||
def forward(self, x):
|
||||
b, c, _, _ = x.size()
|
||||
w = self.pool(x).view(b, c)
|
||||
w = self.fc(w).view(b, c, 1, 1)
|
||||
|
||||
return x * w.expand_as(x)
|
||||
|
||||
|
||||
#------------------------------------------------------------------------------
|
||||
# MODNet Branches
|
||||
#------------------------------------------------------------------------------
|
||||
|
||||
class LRBranch(nn.Module):
|
||||
""" Low Resolution Branch of MODNet
|
||||
"""
|
||||
|
||||
def __init__(self, backbone):
|
||||
super(LRBranch, self).__init__()
|
||||
|
||||
enc_channels = backbone.enc_channels
|
||||
|
||||
self.backbone = backbone
|
||||
self.se_block = SEBlock(enc_channels[4], enc_channels[4], reduction=4)
|
||||
self.conv_lr16x = Conv2dIBNormRelu(enc_channels[4], enc_channels[3], 5, stride=1, padding=2)
|
||||
self.conv_lr8x = Conv2dIBNormRelu(enc_channels[3], enc_channels[2], 5, stride=1, padding=2)
|
||||
self.conv_lr = Conv2dIBNormRelu(enc_channels[2], 1, kernel_size=3, stride=2, padding=1, with_ibn=False, with_relu=False)
|
||||
|
||||
def forward(self, img):
|
||||
enc_features = self.backbone.forward(img)
|
||||
enc2x, enc4x, enc32x = enc_features[0], enc_features[1], enc_features[4]
|
||||
|
||||
enc32x = self.se_block(enc32x)
|
||||
lr16x = F.interpolate(enc32x, scale_factor=2.0, mode='bilinear', align_corners=False)
|
||||
lr16x = self.conv_lr16x(lr16x)
|
||||
lr8x = F.interpolate(lr16x, scale_factor=2.0, mode='bilinear', align_corners=False)
|
||||
lr8x = self.conv_lr8x(lr8x)
|
||||
|
||||
return lr8x, enc2x, enc4x
|
||||
|
||||
|
||||
class HRBranch(nn.Module):
|
||||
""" High Resolution Branch of MODNet
|
||||
"""
|
||||
|
||||
def __init__(self, hr_channels, enc_channels):
|
||||
super(HRBranch, self).__init__()
|
||||
|
||||
self.tohr_enc2x = Conv2dIBNormRelu(enc_channels[0], hr_channels, 1, stride=1, padding=0)
|
||||
self.conv_enc2x = Conv2dIBNormRelu(hr_channels + 3, hr_channels, 3, stride=2, padding=1)
|
||||
|
||||
self.tohr_enc4x = Conv2dIBNormRelu(enc_channels[1], hr_channels, 1, stride=1, padding=0)
|
||||
self.conv_enc4x = Conv2dIBNormRelu(2 * hr_channels, 2 * hr_channels, 3, stride=1, padding=1)
|
||||
|
||||
self.conv_hr4x = nn.Sequential(
|
||||
Conv2dIBNormRelu(3 * hr_channels + 3, 2 * hr_channels, 3, stride=1, padding=1),
|
||||
Conv2dIBNormRelu(2 * hr_channels, 2 * hr_channels, 3, stride=1, padding=1),
|
||||
Conv2dIBNormRelu(2 * hr_channels, hr_channels, 3, stride=1, padding=1),
|
||||
)
|
||||
|
||||
self.conv_hr2x = nn.Sequential(
|
||||
Conv2dIBNormRelu(2 * hr_channels, 2 * hr_channels, 3, stride=1, padding=1),
|
||||
Conv2dIBNormRelu(2 * hr_channels, hr_channels, 3, stride=1, padding=1),
|
||||
Conv2dIBNormRelu(hr_channels, hr_channels, 3, stride=1, padding=1),
|
||||
Conv2dIBNormRelu(hr_channels, hr_channels, 3, stride=1, padding=1),
|
||||
)
|
||||
|
||||
self.conv_hr = nn.Sequential(
|
||||
Conv2dIBNormRelu(hr_channels + 3, hr_channels, 3, stride=1, padding=1),
|
||||
Conv2dIBNormRelu(hr_channels, 1, kernel_size=1, stride=1, padding=0, with_ibn=False, with_relu=False),
|
||||
)
|
||||
|
||||
def forward(self, img, enc2x, enc4x, lr8x):
|
||||
img2x = F.interpolate(img, scale_factor=1/2, mode='bilinear', align_corners=False)
|
||||
img4x = F.interpolate(img, scale_factor=1/4, mode='bilinear', align_corners=False)
|
||||
|
||||
enc2x = self.tohr_enc2x(enc2x)
|
||||
hr4x = self.conv_enc2x(torch.cat((img2x, enc2x), dim=1))
|
||||
|
||||
enc4x = self.tohr_enc4x(enc4x)
|
||||
hr4x = self.conv_enc4x(torch.cat((hr4x, enc4x), dim=1))
|
||||
|
||||
lr4x = F.interpolate(lr8x, scale_factor=2.0, mode='bilinear', align_corners=False)
|
||||
hr4x = self.conv_hr4x(torch.cat((hr4x, lr4x, img4x), dim=1))
|
||||
|
||||
hr2x = F.interpolate(hr4x, scale_factor=2.0, mode='bilinear', align_corners=False)
|
||||
hr2x = self.conv_hr2x(torch.cat((hr2x, enc2x), dim=1))
|
||||
|
||||
return hr2x
|
||||
|
||||
|
||||
class FusionBranch(nn.Module):
|
||||
""" Fusion Branch of MODNet
|
||||
"""
|
||||
|
||||
def __init__(self, hr_channels, enc_channels):
|
||||
super(FusionBranch, self).__init__()
|
||||
self.conv_lr4x = Conv2dIBNormRelu(enc_channels[2], hr_channels, 5, stride=1, padding=2)
|
||||
|
||||
self.conv_f2x = Conv2dIBNormRelu(2 * hr_channels, hr_channels, 3, stride=1, padding=1)
|
||||
self.conv_f = nn.Sequential(
|
||||
Conv2dIBNormRelu(hr_channels + 3, int(hr_channels / 2), 3, stride=1, padding=1),
|
||||
Conv2dIBNormRelu(int(hr_channels / 2), 1, 1, stride=1, padding=0, with_ibn=False, with_relu=False),
|
||||
)
|
||||
|
||||
def forward(self, img, lr8x, hr2x):
|
||||
lr4x = F.interpolate(lr8x, scale_factor=2.0, mode='bilinear', align_corners=False)
|
||||
lr4x = self.conv_lr4x(lr4x)
|
||||
lr2x = F.interpolate(lr4x, scale_factor=2.0, mode='bilinear', align_corners=False)
|
||||
|
||||
f2x = self.conv_f2x(torch.cat((lr2x, hr2x), dim=1))
|
||||
f = F.interpolate(f2x, scale_factor=2.0, mode='bilinear', align_corners=False)
|
||||
f = self.conv_f(torch.cat((f, img), dim=1))
|
||||
pred_matte = torch.sigmoid(f)
|
||||
|
||||
return pred_matte
|
||||
|
||||
|
||||
#------------------------------------------------------------------------------
|
||||
# MODNet
|
||||
#------------------------------------------------------------------------------
|
||||
|
||||
class MODNet(nn.Module):
|
||||
""" Architecture of MODNet
|
||||
"""
|
||||
|
||||
def __init__(self, in_channels=3, hr_channels=32, backbone_arch='mobilenetv2', backbone_pretrained=True):
|
||||
super(MODNet, self).__init__()
|
||||
|
||||
self.in_channels = in_channels
|
||||
self.hr_channels = hr_channels
|
||||
self.backbone_arch = backbone_arch
|
||||
self.backbone_pretrained = backbone_pretrained
|
||||
|
||||
self.backbone = SUPPORTED_BACKBONES[self.backbone_arch](self.in_channels)
|
||||
|
||||
self.lr_branch = LRBranch(self.backbone)
|
||||
self.hr_branch = HRBranch(self.hr_channels, self.backbone.enc_channels)
|
||||
self.f_branch = FusionBranch(self.hr_channels, self.backbone.enc_channels)
|
||||
|
||||
for m in self.modules():
|
||||
if isinstance(m, nn.Conv2d):
|
||||
self._init_conv(m)
|
||||
elif isinstance(m, nn.BatchNorm2d) or isinstance(m, nn.InstanceNorm2d):
|
||||
self._init_norm(m)
|
||||
|
||||
if self.backbone_pretrained:
|
||||
self.backbone.load_pretrained_ckpt()
|
||||
|
||||
def forward(self, img):
|
||||
# NOTE
|
||||
lr_out = self.lr_branch(img)
|
||||
lr8x = lr_out[0]
|
||||
enc2x = lr_out[1]
|
||||
enc4x = lr_out[2]
|
||||
|
||||
hr2x = self.hr_branch(img, enc2x, enc4x, lr8x)
|
||||
|
||||
pred_matte = self.f_branch(img, lr8x, hr2x)
|
||||
|
||||
return pred_matte
|
||||
|
||||
def freeze_norm(self):
|
||||
norm_types = [nn.BatchNorm2d, nn.InstanceNorm2d]
|
||||
for m in self.modules():
|
||||
for n in norm_types:
|
||||
if isinstance(m, n):
|
||||
m.eval()
|
||||
continue
|
||||
|
||||
def _init_conv(self, conv):
|
||||
nn.init.kaiming_uniform_(
|
||||
conv.weight, a=0, mode='fan_in', nonlinearity='relu')
|
||||
if conv.bias is not None:
|
||||
nn.init.constant_(conv.bias, 0)
|
||||
|
||||
def _init_norm(self, norm):
|
||||
if norm.weight is not None:
|
||||
nn.init.constant_(norm.weight, 1)
|
||||
nn.init.constant_(norm.bias, 0)
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,147 @@
|
||||
#!/usr/bin/env python3
|
||||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
自动修复 subtitle_cover_generator_simple.py 的脚本
|
||||
修复人物描边功能和添加蒙版功能
|
||||
"""
|
||||
|
||||
import sys
|
||||
import re
|
||||
import io
|
||||
|
||||
# 设置标准输出为UTF-8
|
||||
sys.stdout = io.TextIOWrapper(sys.stdout.buffer, encoding='utf-8')
|
||||
sys.stderr = io.TextIOWrapper(sys.stderr.buffer, encoding='utf-8')
|
||||
|
||||
def fix_file():
|
||||
# 读取原文件
|
||||
with open('subtitle_cover_generator_simple.py', 'r', encoding='utf-8') as f:
|
||||
content = f.read()
|
||||
|
||||
# 检查是否已经修复过
|
||||
if '# 获取alpha通道作为mask' in content and 'mask_config = template.get' in content:
|
||||
print('文件已经修复过,无需再次修复')
|
||||
return
|
||||
|
||||
# 修复1: 替换人物描边的TODO部分
|
||||
old_stroke_code = ''' # 合成带描边的人像
|
||||
if portrait.mode == 'RGBA':
|
||||
result = Image.new('RGB', (width, height))
|
||||
result.paste(image, (0, 0))
|
||||
# TODO: 实现描边效果（需要更复杂的图像处理）
|
||||
result.paste(portrait, (0, 0), portrait)
|
||||
image = result
|
||||
except Exception as e:
|
||||
print(f"Portrait stroke failed: {e}", file=sys.stderr)'''
|
||||
|
||||
new_stroke_code = ''' # 合成带描边的人像
|
||||
if portrait.mode == 'RGBA':
|
||||
# 获取alpha通道作为mask
|
||||
alpha = portrait.split()[3]
|
||||
|
||||
# 对mask进行边缘检测
|
||||
alpha_np = np.array(alpha)
|
||||
edges = cv2.Canny(alpha_np, 30, 100)
|
||||
|
||||
# 转换描边颜色
|
||||
stroke_rgb = self._hex_to_rgb(stroke_color)
|
||||
outline_layer = Image.new('RGBA', (width, height), (*stroke_rgb, 0))
|
||||
|
||||
# 创建多层描边
|
||||
for i in range(stroke_width, 0, -1):
|
||||
# 外层更粗的描边
|
||||
kernel_size = i * 3
|
||||
kernel = np.ones((kernel_size, kernel_size), np.uint8)
|
||||
layer_edges = cv2.dilate(edges, kernel, iterations=1)
|
||||
|
||||
# 转换为PIL图像
|
||||
layer_pil = Image.fromarray(layer_edges)
|
||||
layer_colored = Image.new('RGBA', (width, height), (*stroke_rgb, 255))
|
||||
layer_colored.putalpha(layer_pil)
|
||||
|
||||
# 合并到描边层
|
||||
outline_layer = Image.alpha_composite(outline_layer, layer_colored)
|
||||
|
||||
# 创建结果：背景 + 描边 + 人像
|
||||
result = Image.new('RGB', (width, height))
|
||||
result.paste(image, (0, 0))
|
||||
|
||||
# 将描边和人像合成到背景上
|
||||
result_rgba = result.convert('RGBA')
|
||||
result_rgba = Image.alpha_composite(result_rgba, outline_layer)
|
||||
result_rgba = Image.alpha_composite(result_rgba, portrait)
|
||||
image = result_rgba.convert('RGB')
|
||||
|
||||
print("Portrait stroke effect applied successfully", file=sys.stderr)
|
||||
except Exception as e:
|
||||
print(f"Portrait stroke failed: {e}", file=sys.stderr)
|
||||
import traceback
|
||||
traceback.print_exc(file=sys.stderr)'''
|
||||
|
||||
if old_stroke_code in content:
|
||||
content = content.replace(old_stroke_code, new_stroke_code)
|
||||
print('✓ 已修复人物描边功能')
|
||||
else:
|
||||
print('✗ 未找到人物描边TODO代码,可能已被修改')
|
||||
|
||||
# 修复2: 在"# 3. 添加标题"之前插入蒙版代码
|
||||
mask_code = '''
|
||||
# 3. 添加蒙版（如果有）
|
||||
mask_config = template.get('mask', {})
|
||||
mask_path = mask_config.get('imagePath', '')
|
||||
if mask_path and os.path.exists(mask_path):
|
||||
try:
|
||||
print(f"Adding mask from: {mask_path}", file=sys.stderr)
|
||||
mask_img = Image.open(mask_path).convert('RGBA')
|
||||
|
||||
# 获取蒙版位置和大小
|
||||
mask_pos = mask_config.get('position', {'x': 10, 'y': 10})
|
||||
mask_size = mask_config.get('size', {'width': 30, 'height': 20})
|
||||
mask_opacity = mask_config.get('opacity', 0.8)
|
||||
|
||||
mask_width = int(mask_size['width'] / 100 * width)
|
||||
mask_height = int(mask_size['height'] / 100 * height)
|
||||
mask_x = int(mask_pos['x'] / 100 * width)
|
||||
mask_y = int(mask_pos['y'] / 100 * height)
|
||||
|
||||
# 调整蒙版大小和透明度
|
||||
mask_img = mask_img.resize((mask_width, mask_height), Image.LANCZOS)
|
||||
if mask_opacity < 1.0:
|
||||
alpha = mask_img.split()[3]
|
||||
alpha = alpha.point(lambda p: int(p * mask_opacity))
|
||||
mask_img.putalpha(alpha)
|
||||
|
||||
# 将蒙版粘贴到图像上
|
||||
image_rgba = image.convert('RGBA')
|
||||
image_rgba.paste(mask_img, (mask_x - mask_width // 2, mask_y - mask_height // 2), mask_img)
|
||||
image = image_rgba.convert('RGB')
|
||||
|
||||
print("Mask applied successfully", file=sys.stderr)
|
||||
except Exception as e:
|
||||
print(f"Failed to add mask: {e}", file=sys.stderr)
|
||||
import traceback
|
||||
traceback.print_exc(file=sys.stderr)
|
||||
|
||||
# 4. 添加标题'''
|
||||
|
||||
# 查找并替换 "# 3. 添加标题"
|
||||
if ' # 3. 添加标题' in content:
|
||||
content = content.replace(' # 3. 添加标题', mask_code)
|
||||
print('✓ 已添加蒙版功能')
|
||||
else:
|
||||
print('✗ 未找到"# 3. 添加标题"标记')
|
||||
|
||||
# 写回文件
|
||||
with open('subtitle_cover_generator_simple.py', 'w', encoding='utf-8') as f:
|
||||
f.write(content)
|
||||
|
||||
print('\n修复完成! 已创建备份文件: subtitle_cover_generator_simple.py.backup')
|
||||
|
||||
if __name__ == '__main__':
|
||||
try:
|
||||
fix_file()
|
||||
except Exception as e:
|
||||
print(f'修复失败: {e}')
|
||||
import traceback
|
||||
traceback.print_exc()
|
||||
sys.exit(1)
|
||||
@@ -0,0 +1,361 @@
|
||||
#!/usr/bin/env python
|
||||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
视频封面生成器
|
||||
支持从视频提取帧并添加标题文字
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
from PIL import Image, ImageDraw, ImageFont, ImageFilter
|
||||
import cv2
|
||||
|
||||
# ⚠️ 修复模块路径：确保 Python 能找到 modules 和其他本地模块
|
||||
# 将脚本所在目录（python 目录）添加到 sys.path
|
||||
script_dir = os.path.dirname(os.path.abspath(__file__))
|
||||
if script_dir not in sys.path:
|
||||
sys.path.insert(0, script_dir)
|
||||
print(f"[PATH] Added to sys.path: {script_dir}", file=sys.stderr)
|
||||
|
||||
def extract_frame_from_video(video_path, timestamp=0):
|
||||
"""从视频中提取指定时间的帧"""
|
||||
try:
|
||||
cap = cv2.VideoCapture(video_path)
|
||||
|
||||
# 设置到指定时间（秒）
|
||||
cap.set(cv2.CAP_PROP_POS_MSEC, timestamp * 1000)
|
||||
|
||||
ret, frame = cap.read()
|
||||
cap.release()
|
||||
|
||||
if not ret:
|
||||
return None
|
||||
|
||||
# 转换BGR到RGB
|
||||
frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
|
||||
return Image.fromarray(frame_rgb)
|
||||
|
||||
except Exception as e:
|
||||
print(f"Error extracting frame: {e}", file=sys.stderr)
|
||||
return None
|
||||
|
||||
def hex_to_rgb(hex_color):
|
||||
"""将十六进制颜色转换为RGB元组"""
|
||||
if isinstance(hex_color, str):
|
||||
hex_color = hex_color.lstrip('#')
|
||||
return tuple(int(hex_color[i:i+2], 16) for i in (0, 2, 4))
|
||||
return tuple(hex_color)
|
||||
|
||||
def add_text_to_image(img, title, config):
|
||||
"""在图像上添加标题文字"""
|
||||
pil_img = img.copy().convert('RGBA')
|
||||
|
||||
width, height = pil_img.size
|
||||
|
||||
# 获取配置参数 - 支持新格式（来自前端的配置）
|
||||
# 新格式：titleFontSize, titleColor, titleFontFamily, titlePosition, titleStrokeWidth, titleStrokeColor
|
||||
# 旧格式：fontSize, fontColor, fontPath, position, style
|
||||
|
||||
font_size = config.get('titleFontSize', config.get('fontSize', int(height * 0.08)))
|
||||
|
||||
# 处理颜色格式：新格式是十六进制字符串（如 "#eb1414"），旧格式是RGB元组
|
||||
title_color = config.get('titleColor', config.get('fontColor', '#FFFFFF'))
|
||||
font_color = hex_to_rgb(title_color)
|
||||
|
||||
# 阴影参数
|
||||
shadow_color = hex_to_rgb(config.get('titleShadowColor', '#000000'))
|
||||
shadow_offset_x = config.get('titleShadowOffsetX', 0)
|
||||
shadow_offset_y = config.get('titleShadowOffsetY', 0)
|
||||
shadow_blur_raw = config.get('titleShadowBlur', 0)
|
||||
# 修复：PIL的GaussianBlur效果比CSS重，需要缩小系数以匹配前端预览
|
||||
# 经验值：CSS box-shadow blur 与 PIL GaussianBlur 的比例约为 2.5:1
|
||||
shadow_blur = int(shadow_blur_raw * 0.4) if shadow_blur_raw > 0 else 0
|
||||
|
||||
# 描边参数
|
||||
stroke_width = config.get('titleStrokeWidth', 0)
|
||||
stroke_color = hex_to_rgb(config.get('titleStrokeColor', '#000000'))
|
||||
|
||||
print(f"[DEBUG] Text rendering params: stroke_width={stroke_width}, stroke_color={stroke_color}", file=sys.stderr)
|
||||
|
||||
font_family = config.get('titleFontFamily', config.get('fontPath', 'NotoSerifCJK-VF'))
|
||||
|
||||
# 使用 font_manager 查找字体（支持系统字体和ziti目录字体）
|
||||
font_path = None
|
||||
try:
|
||||
from modules.font_manager import get_font_manager
|
||||
font_manager = get_font_manager()
|
||||
font_path = font_manager.find_font(font_family, 400)
|
||||
if font_path:
|
||||
print(f"[DEBUG] Font found via font_manager: {font_family} -> {font_path}", file=sys.stderr)
|
||||
else:
|
||||
print(f"[DEBUG] Font not found via font_manager: {font_family}, using fallback", file=sys.stderr)
|
||||
except Exception as e:
|
||||
print(f"[DEBUG] font_manager not available: {e}, using hardcoded fonts", file=sys.stderr)
|
||||
|
||||
if not font_path:
|
||||
raise RuntimeError(f"Bundled font not found: {font_family}")
|
||||
|
||||
|
||||
# 位置：新格式是百分比坐标 {x, y}，旧格式是 'top', 'center', 'bottom'
|
||||
title_position = config.get('titlePosition', {})
|
||||
if isinstance(title_position, dict):
|
||||
position = 'custom'
|
||||
else:
|
||||
position = config.get('position', 'top') # top, center, bottom
|
||||
|
||||
style = config.get('style', 'default') # default, outline, blur_bg, gradient, split
|
||||
|
||||
# 加载字体
|
||||
try:
|
||||
font = ImageFont.truetype(font_path, font_size)
|
||||
except:
|
||||
font = ImageFont.load_default()
|
||||
|
||||
# 创建绘图对象
|
||||
draw = ImageDraw.Draw(pil_img)
|
||||
|
||||
# 计算文字尺寸
|
||||
bbox = draw.textbbox((0, 0), title, font=font)
|
||||
text_width = bbox[2] - bbox[0]
|
||||
text_height = bbox[3] - bbox[1]
|
||||
|
||||
# 计算文字位置
|
||||
if position == 'custom' and isinstance(title_position, dict):
|
||||
# 使用自定义位置（百分比）
|
||||
x = int(width * title_position.get('x', 50) / 100) - text_width // 2
|
||||
y = int(height * title_position.get('y', 80) / 100) - text_height // 2
|
||||
else:
|
||||
x = (width - text_width) // 2
|
||||
if position == 'top':
|
||||
y = int(height * 0.08)
|
||||
elif position == 'center':
|
||||
y = (height - text_height) // 2
|
||||
elif position == 'bottom':
|
||||
y = int(height * 0.85) - text_height
|
||||
else:
|
||||
y = int(height * 0.08)
|
||||
|
||||
# 辅助函数：绘制带阴影和描边的文字
|
||||
def draw_text_with_effects(draw_obj, pos_x, pos_y, text, font, color, stroke_w, stroke_c, shadow_ox, shadow_oy, shadow_b, shadow_c):
|
||||
"""绘制带阴影和描边的文字"""
|
||||
# 1. 先绘制阴影（如果有偏移）
|
||||
if shadow_ox != 0 or shadow_oy != 0:
|
||||
if shadow_b > 0:
|
||||
# 创建阴影图层并模糊
|
||||
shadow_layer = Image.new('RGBA', pil_img.size, (0, 0, 0, 0))
|
||||
shadow_draw = ImageDraw.Draw(shadow_layer)
|
||||
shadow_draw.text((pos_x + shadow_ox, pos_y + shadow_oy), text, font=font, fill=shadow_c + (180,))
|
||||
shadow_layer = shadow_layer.filter(ImageFilter.GaussianBlur(radius=shadow_b))
|
||||
return shadow_layer
|
||||
else:
|
||||
draw_obj.text((pos_x + shadow_ox, pos_y + shadow_oy), text, font=font, fill=shadow_c + (180,))
|
||||
|
||||
# 2. 绘制描边
|
||||
if stroke_w > 0:
|
||||
for adj_x in range(-stroke_w, stroke_w + 1):
|
||||
for adj_y in range(-stroke_w, stroke_w + 1):
|
||||
if adj_x != 0 or adj_y != 0:
|
||||
draw_obj.text((pos_x + adj_x, pos_y + adj_y), text, font=font, fill=stroke_c + (255,))
|
||||
|
||||
# 3. 绘制主文字
|
||||
draw_obj.text((pos_x, pos_y), text, font=font, fill=color + (255,))
|
||||
return None
|
||||
|
||||
# 应用样式
|
||||
if style == 'default':
|
||||
# 默认风格：黑色半透明背景条
|
||||
padding = 30
|
||||
overlay = Image.new('RGBA', pil_img.size, (0, 0, 0, 0))
|
||||
overlay_draw = ImageDraw.Draw(overlay)
|
||||
overlay_draw.rectangle(
|
||||
[x - padding, y - padding, x + text_width + padding, y + text_height + padding],
|
||||
fill=(0, 0, 0, 180)
|
||||
)
|
||||
pil_img = Image.alpha_composite(pil_img, overlay)
|
||||
draw = ImageDraw.Draw(pil_img)
|
||||
|
||||
# 绘制阴影
|
||||
if shadow_offset_x != 0 or shadow_offset_y != 0:
|
||||
if shadow_blur > 0:
|
||||
shadow_layer = Image.new('RGBA', pil_img.size, (0, 0, 0, 0))
|
||||
shadow_draw = ImageDraw.Draw(shadow_layer)
|
||||
shadow_draw.text((x + shadow_offset_x, y + shadow_offset_y), title, font=font, fill=shadow_color + (180,))
|
||||
shadow_layer = shadow_layer.filter(ImageFilter.GaussianBlur(radius=shadow_blur))
|
||||
pil_img = Image.alpha_composite(pil_img, shadow_layer)
|
||||
draw = ImageDraw.Draw(pil_img)
|
||||
else:
|
||||
draw.text((x + shadow_offset_x, y + shadow_offset_y), title, font=font, fill=shadow_color + (180,))
|
||||
|
||||
# 绘制描边
|
||||
if stroke_width > 0:
|
||||
print(f"[DEBUG] Drawing text stroke: width={stroke_width}, color={stroke_color}", file=sys.stderr)
|
||||
for adj_x in range(-stroke_width, stroke_width + 1):
|
||||
for adj_y in range(-stroke_width, stroke_width + 1):
|
||||
if adj_x != 0 or adj_y != 0:
|
||||
draw.text((x + adj_x, y + adj_y), title, font=font, fill=stroke_color + (255,))
|
||||
else:
|
||||
print(f"[DEBUG] Stroke disabled: stroke_width={stroke_width}", file=sys.stderr)
|
||||
|
||||
draw.text((x, y), title, font=font, fill=font_color + (255,))
|
||||
|
||||
elif style == 'outline':
|
||||
# 描边风格
|
||||
# 绘制阴影
|
||||
if shadow_offset_x != 0 or shadow_offset_y != 0:
|
||||
if shadow_blur > 0:
|
||||
shadow_layer = Image.new('RGBA', pil_img.size, (0, 0, 0, 0))
|
||||
shadow_draw = ImageDraw.Draw(shadow_layer)
|
||||
shadow_draw.text((x + shadow_offset_x, y + shadow_offset_y), title, font=font, fill=shadow_color + (180,))
|
||||
shadow_layer = shadow_layer.filter(ImageFilter.GaussianBlur(radius=shadow_blur))
|
||||
pil_img = Image.alpha_composite(pil_img, shadow_layer)
|
||||
draw = ImageDraw.Draw(pil_img)
|
||||
else:
|
||||
draw.text((x + shadow_offset_x, y + shadow_offset_y), title, font=font, fill=shadow_color + (180,))
|
||||
|
||||
outline_width = stroke_width if stroke_width > 0 else 3
|
||||
outline_color = stroke_color
|
||||
for adj_x in range(-outline_width, outline_width + 1):
|
||||
for adj_y in range(-outline_width, outline_width + 1):
|
||||
draw.text((x + adj_x, y + adj_y), title, font=font, fill=outline_color + (255,))
|
||||
draw.text((x, y), title, font=font, fill=font_color + (255,))
|
||||
|
||||
elif style == 'blur_bg':
|
||||
# 模糊背景风格
|
||||
padding = 50
|
||||
bg_region = pil_img.crop((
|
||||
max(0, x - padding),
|
||||
max(0, y - padding),
|
||||
min(width, x + text_width + padding),
|
||||
min(height, y + text_height + padding)
|
||||
))
|
||||
bg_region = bg_region.filter(ImageFilter.GaussianBlur(radius=15))
|
||||
pil_img.paste(bg_region, (max(0, x - padding), max(0, y - padding)))
|
||||
|
||||
# 添加半透明遮罩
|
||||
overlay = Image.new('RGBA', pil_img.size, (0, 0, 0, 0))
|
||||
overlay_draw = ImageDraw.Draw(overlay)
|
||||
overlay_draw.rectangle(
|
||||
[x - padding, y - padding, x + text_width + padding, y + text_height + padding],
|
||||
fill=(0, 0, 0, 120)
|
||||
)
|
||||
pil_img = Image.alpha_composite(pil_img, overlay)
|
||||
draw = ImageDraw.Draw(pil_img)
|
||||
draw.text((x, y), title, font=font, fill=font_color + (255,))
|
||||
|
||||
elif style == 'gradient':
|
||||
# 渐变背景风格
|
||||
gradient_height = text_height + 100
|
||||
gradient = Image.new('RGBA', (width, gradient_height), (0, 0, 0, 0))
|
||||
gradient_draw = ImageDraw.Draw(gradient)
|
||||
for i in range(gradient_height):
|
||||
alpha = int(200 * (i / gradient_height))
|
||||
gradient_draw.rectangle([0, i, width, i + 1], fill=(0, 0, 0, alpha))
|
||||
|
||||
pil_img.paste(gradient, (0, max(0, y - 50)), gradient)
|
||||
draw = ImageDraw.Draw(pil_img)
|
||||
draw.text((x, y), title, font=font, fill=font_color + (255,))
|
||||
|
||||
elif style == 'split':
|
||||
# 分栏风格
|
||||
padding = 40
|
||||
left_x = padding
|
||||
|
||||
# 左侧标题背景
|
||||
overlay = Image.new('RGBA', pil_img.size, (0, 0, 0, 0))
|
||||
overlay_draw = ImageDraw.Draw(overlay)
|
||||
overlay_draw.rectangle(
|
||||
[left_x - 20, y - 20, left_x + text_width + 20, y + text_height + 20],
|
||||
fill=(0, 0, 0, 180)
|
||||
)
|
||||
pil_img = Image.alpha_composite(pil_img, overlay)
|
||||
draw = ImageDraw.Draw(pil_img)
|
||||
draw.text((left_x, y), title, font=font, fill=font_color + (255,))
|
||||
|
||||
# 右侧装饰
|
||||
accent_text = config.get('accentText', '✨')
|
||||
accent_bbox = draw.textbbox((0, 0), accent_text, font=font)
|
||||
accent_width = accent_bbox[2] - accent_bbox[0]
|
||||
right_x = width - accent_width - padding
|
||||
|
||||
overlay = Image.new('RGBA', pil_img.size, (0, 0, 0, 0))
|
||||
overlay_draw = ImageDraw.Draw(overlay)
|
||||
overlay_draw.rectangle(
|
||||
[right_x - 25, y - 25, right_x + accent_width + 25, y + text_height + 25],
|
||||
fill=(255, 193, 7, 200)
|
||||
)
|
||||
pil_img = Image.alpha_composite(pil_img, overlay)
|
||||
draw = ImageDraw.Draw(pil_img)
|
||||
draw.text((right_x, y), accent_text, font=font, fill=(0, 0, 0, 255))
|
||||
|
||||
return pil_img.convert('RGB')
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description='Generate video cover with title')
|
||||
parser.add_argument('--video', required=True, help='Path to video file')
|
||||
parser.add_argument('--title', required=True, help='Title text')
|
||||
parser.add_argument('--output', required=True, help='Output image path')
|
||||
parser.add_argument('--config', default='{}', help='JSON configuration')
|
||||
parser.add_argument('--preview', action='store_true', help='Preview mode')
|
||||
parser.add_argument('--timestamp', type=float, default=0, help='Frame timestamp in seconds')
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
try:
|
||||
# 解析配置
|
||||
config = json.loads(args.config)
|
||||
|
||||
# 诊断日志：输出接收到的配置参数
|
||||
print(f"[DEBUG] Received config: {json.dumps(config, indent=2, ensure_ascii=False)}", file=sys.stderr)
|
||||
print(f"[DEBUG] titleStrokeWidth: {config.get('titleStrokeWidth', 'NOT SET')}", file=sys.stderr)
|
||||
print(f"[DEBUG] titleStrokeColor: {config.get('titleStrokeColor', 'NOT SET')}", file=sys.stderr)
|
||||
print(f"[DEBUG] titleFontFamily: {config.get('titleFontFamily', 'NOT SET')}", file=sys.stderr)
|
||||
|
||||
# 检查输入是图像还是视频
|
||||
video_path = args.video
|
||||
if video_path.lower().endswith(('.jpg', '.jpeg', '.png', '.bmp', '.webp')):
|
||||
# 直接加载图像
|
||||
print(f"Loading image: {video_path}")
|
||||
try:
|
||||
frame = Image.open(video_path)
|
||||
except Exception as e:
|
||||
print(f"Error: Failed to load image: {e}", file=sys.stderr)
|
||||
sys.exit(1)
|
||||
else:
|
||||
# 提取视频帧
|
||||
print(f"Extracting frame from video: {video_path}")
|
||||
frame = extract_frame_from_video(video_path, args.timestamp)
|
||||
|
||||
if frame is None:
|
||||
print("Error: Failed to extract frame from video", file=sys.stderr)
|
||||
sys.exit(1)
|
||||
|
||||
# 调整图片大小
|
||||
width, height = frame.size
|
||||
max_width = config.get('maxWidth', 1920)
|
||||
if width > max_width:
|
||||
scale = max_width / width
|
||||
new_width = max_width
|
||||
new_height = int(height * scale)
|
||||
frame = frame.resize((new_width, new_height), Image.Resampling.LANCZOS)
|
||||
|
||||
# 添加标题
|
||||
print(f"Adding title: {args.title}")
|
||||
result = add_text_to_image(frame, args.title, config)
|
||||
|
||||
# 保存结果
|
||||
os.makedirs(os.path.dirname(args.output), exist_ok=True)
|
||||
result.save(args.output, quality=95)
|
||||
print(f"Cover saved to: {args.output}")
|
||||
|
||||
sys.exit(0)
|
||||
|
||||
except Exception as e:
|
||||
print(f"Error: {e}", file=sys.stderr)
|
||||
import traceback
|
||||
traceback.print_exc()
|
||||
sys.exit(1)
|
||||
|
||||
if __name__ == '__main__':
|
||||
main()
|
||||
@@ -0,0 +1,221 @@
|
||||
#!/usr/bin/env python3
|
||||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
实时封面预览生成器
|
||||
Generate Real-time Cover Previews
|
||||
"""
|
||||
|
||||
import os
|
||||
import sys
|
||||
import json
|
||||
import hashlib
|
||||
import tempfile
|
||||
from pathlib import Path
|
||||
from typing import Dict, List, Any, Optional
|
||||
import subprocess
|
||||
import threading
|
||||
import time
|
||||
|
||||
# 导入封面生成器
|
||||
from subtitle_cover_generator_simple import SubtitleCoverGenerator
|
||||
|
||||
class CoverPreviewGenerator:
|
||||
"""实时封面预览生成器"""
|
||||
|
||||
def __init__(self):
|
||||
self.temp_dir = os.path.join(tempfile.gettempdir(), 'cover_previews')
|
||||
os.makedirs(self.temp_dir, exist_ok=True)
|
||||
|
||||
# 缓存已生成的预览，避免重复生成
|
||||
self.preview_cache = {}
|
||||
|
||||
# 定义封面样式
|
||||
self.COVER_STYLES = [
|
||||
'default',
|
||||
'blur-bg',
|
||||
'outline',
|
||||
'gradient',
|
||||
'split'
|
||||
]
|
||||
|
||||
def get_video_hash(self, video_path: str) -> str:
|
||||
"""获取视频文件的哈希值用于缓存"""
|
||||
hash_md5 = hashlib.md5()
|
||||
with open(video_path, "rb") as f:
|
||||
for chunk in iter(lambda: f.read(4096), b""):
|
||||
hash_md5.update(chunk)
|
||||
return hash_md5.hexdigest()[:16]
|
||||
|
||||
def extract_first_frame(self, video_path: str, output_path: str) -> bool:
|
||||
"""从视频中提取第一帧"""
|
||||
try:
|
||||
# 使用ffmpeg提取第一帧
|
||||
cmd = [
|
||||
'ffmpeg', '-i', video_path,
|
||||
'-vf', 'select=eq(n\\,0)', # 选择第一帧
|
||||
'-q:v', '2', # 高质量
|
||||
'-f', 'image2',
|
||||
'-y', # 覆盖输出
|
||||
output_path
|
||||
]
|
||||
|
||||
result = subprocess.run(cmd, capture_output=True, text=True, timeout=30)
|
||||
return result.returncode == 0 and os.path.exists(output_path)
|
||||
|
||||
except Exception as e:
|
||||
print(f"Error extracting frame: {e}", file=sys.stderr)
|
||||
return False
|
||||
|
||||
def generate_preview(self, video_path: str, style: str, output_path: str) -> bool:
|
||||
"""生成单个样式的封面预览"""
|
||||
try:
|
||||
# 创建临时帧文件
|
||||
frame_path = os.path.join(self.temp_dir, f"frame_{os.path.basename(video_path)}.jpg")
|
||||
|
||||
# 提取第一帧
|
||||
if not self.extract_first_frame(video_path, frame_path):
|
||||
print(f"Failed to extract frame from {video_path}", file=sys.stderr)
|
||||
return False
|
||||
|
||||
# 使用封面生成器创建预览
|
||||
config = {
|
||||
'videoPath': video_path, # 保持原视频路径，用于可能的其他处理
|
||||
'scriptText': '实时预览封面效果', # 简短预览文本
|
||||
'subtitleStyle': 'default',
|
||||
'coverStyle': style,
|
||||
'outputDir': os.path.dirname(output_path),
|
||||
'coverOnly': True,
|
||||
'useExtractedFrame': True, # 使用已提取的帧
|
||||
'extractedFramePath': frame_path
|
||||
}
|
||||
|
||||
generator = SubtitleCoverGenerator(config)
|
||||
cover_path = generator.generate_cover()
|
||||
|
||||
if cover_path and os.path.exists(cover_path):
|
||||
# 移动到最终位置
|
||||
os.rename(cover_path, output_path)
|
||||
print(f"Generated preview: {output_path}", file=sys.stderr)
|
||||
return True
|
||||
else:
|
||||
print(f"Failed to generate cover for style: {style}", file=sys.stderr)
|
||||
return False
|
||||
|
||||
except Exception as e:
|
||||
print(f"Error generating preview for {style}: {e}", file=sys.stderr)
|
||||
return False
|
||||
finally:
|
||||
# 清理临时文件
|
||||
if os.path.exists(frame_path):
|
||||
try:
|
||||
os.remove(frame_path)
|
||||
except:
|
||||
pass
|
||||
|
||||
def generate_all_previews(self, video_path: str) -> Dict[str, Any]:
|
||||
"""生成所有样式的预览"""
|
||||
if not os.path.exists(video_path):
|
||||
return {"error": f"Video file not found: {video_path}"}
|
||||
|
||||
video_hash = self.get_video_hash(video_path)
|
||||
|
||||
# 检查缓存
|
||||
if video_hash in self.preview_cache:
|
||||
print(f"Using cached previews for {video_hash}", file=sys.stderr)
|
||||
return self.preview_cache[video_hash]
|
||||
|
||||
previews = []
|
||||
errors = []
|
||||
|
||||
print(f"Generating {len(self.COVER_STYLES)} previews for: {video_path}", file=sys.stderr)
|
||||
|
||||
# 并行生成所有预览
|
||||
threads = []
|
||||
results = {}
|
||||
|
||||
def generate_single_preview(style):
|
||||
output_path = os.path.join(self.temp_dir, f"preview_{video_hash}_{style}.png")
|
||||
success = self.generate_preview(video_path, style, output_path)
|
||||
results[style] = {
|
||||
'style': style,
|
||||
'path': output_path if success else None,
|
||||
'success': success
|
||||
}
|
||||
|
||||
# 启动所有预览生成线程
|
||||
for style in self.COVER_STYLES:
|
||||
thread = threading.Thread(target=generate_single_preview, args=(style,))
|
||||
thread.start()
|
||||
threads.append(thread)
|
||||
|
||||
# 等待所有线程完成
|
||||
for thread in threads:
|
||||
thread.join()
|
||||
|
||||
# 整理结果
|
||||
for style in self.COVER_STYLES:
|
||||
result = results.get(style, {})
|
||||
if result.get('success') and result.get('path'):
|
||||
previews.append({
|
||||
'style': style,
|
||||
'url': f"/temp/cover_previews/{os.path.basename(result['path'])}"
|
||||
})
|
||||
else:
|
||||
errors.append(style)
|
||||
|
||||
response = {
|
||||
'previews': previews,
|
||||
'videoHash': video_hash,
|
||||
'tempDir': self.temp_dir
|
||||
}
|
||||
|
||||
if errors:
|
||||
response['errors'] = errors
|
||||
|
||||
# 缓存结果
|
||||
self.preview_cache[video_hash] = response
|
||||
|
||||
print(f"Generated {len(previews)} previews, {len(errors)} failed", file=sys.stderr)
|
||||
return response
|
||||
|
||||
def cleanup_old_previews(self, max_age_hours: int = 24):
|
||||
"""清理旧的预览文件"""
|
||||
try:
|
||||
current_time = time.time()
|
||||
for filename in os.listdir(self.temp_dir):
|
||||
filepath = os.path.join(self.temp_dir, filename)
|
||||
if os.path.isfile(filepath):
|
||||
file_age = current_time - os.path.getmtime(filepath)
|
||||
if file_age > max_age_hours * 3600:
|
||||
os.remove(filepath)
|
||||
print(f"Cleaned up old preview: {filename}", file=sys.stderr)
|
||||
except Exception as e:
|
||||
print(f"Error during cleanup: {e}", file=sys.stderr)
|
||||
|
||||
def main():
|
||||
"""命令行接口"""
|
||||
import argparse
|
||||
|
||||
parser = argparse.ArgumentParser(description='Generate real-time cover previews')
|
||||
parser.add_argument('video_path', help='Path to the video file')
|
||||
parser.add_argument('--output-dir', help='Output directory for previews')
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
if not os.path.exists(args.video_path):
|
||||
print(f"Error: Video file not found: {args.video_path}", file=sys.stderr)
|
||||
sys.exit(1)
|
||||
|
||||
generator = CoverPreviewGenerator()
|
||||
|
||||
if args.output_dir:
|
||||
generator.temp_dir = args.output_dir
|
||||
os.makedirs(generator.temp_dir, exist_ok=True)
|
||||
|
||||
result = generator.generate_all_previews(args.video_path)
|
||||
|
||||
# 输出JSON结果
|
||||
print(json.dumps(result, ensure_ascii=False))
|
||||
|
||||
if __name__ == '__main__':
|
||||
main()
|
||||
@@ -0,0 +1,235 @@
|
||||
#!/usr/bin/env python3
|
||||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
自定义模板封面修复补丁
|
||||
将这段代码复制到 subtitle_cover_generator_simple.py 替换 _generate_custom_template_cover 方法
|
||||
"""
|
||||
|
||||
def _generate_custom_template_cover(self, frame):
|
||||
"""
|
||||
使用自定义模板生成封面 (修复版)
|
||||
|
||||
Args:
|
||||
frame: 视频帧
|
||||
|
||||
Returns:
|
||||
生成的封面路径
|
||||
"""
|
||||
try:
|
||||
print("Generating custom template cover (fixed version)...", file=sys.stderr)
|
||||
template = self.template
|
||||
|
||||
# 转换为PIL图像
|
||||
image = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
|
||||
width, height = image.size
|
||||
|
||||
# 1. 应用背景模糊
|
||||
if template.get('background', {}).get('blurEnabled', False):
|
||||
blur_radius = template['background'].get('blurRadius', 15)
|
||||
print(f"Applying background blur: {blur_radius}", file=sys.stderr)
|
||||
image = image.filter(ImageFilter.GaussianBlur(radius=blur_radius))
|
||||
|
||||
# 2. 处理人像描边（如果启用RMBG）
|
||||
if template.get('portrait', {}).get('strokeEnabled', False):
|
||||
if RMBG_AVAILABLE:
|
||||
try:
|
||||
print("Extracting portrait for stroke effect...", file=sys.stderr)
|
||||
portrait = rembg_remove(image)
|
||||
|
||||
# 创建描边效果
|
||||
stroke_color = template['portrait'].get('strokeColor', '#FFD700')
|
||||
stroke_width = template['portrait'].get('strokeWidth', 3)
|
||||
stroke_type = template['portrait'].get('strokeType', 'solid')
|
||||
|
||||
print(f"Stroke: color={stroke_color}, width={stroke_width}, type={stroke_type}", file=sys.stderr)
|
||||
|
||||
# 合成带描边的人像
|
||||
if portrait.mode == 'RGBA':
|
||||
# 获取alpha通道作为mask
|
||||
alpha = portrait.split()[3]
|
||||
|
||||
# 对mask进行边缘检测
|
||||
alpha_np = np.array(alpha)
|
||||
edges = cv2.Canny(alpha_np, 30, 100)
|
||||
|
||||
# 转换描边颜色
|
||||
stroke_rgb = self._hex_to_rgb(stroke_color)
|
||||
outline_layer = Image.new('RGBA', (width, height), (*stroke_rgb, 0))
|
||||
|
||||
# 创建多层描边
|
||||
for i in range(stroke_width, 0, -1):
|
||||
# 外层更粗的描边
|
||||
kernel_size = i * 3
|
||||
kernel = np.ones((kernel_size, kernel_size), np.uint8)
|
||||
layer_edges = cv2.dilate(edges, kernel, iterations=1)
|
||||
|
||||
# 转换为PIL图像
|
||||
layer_pil = Image.fromarray(layer_edges)
|
||||
layer_colored = Image.new('RGBA', (width, height), (*stroke_rgb, 255))
|
||||
layer_colored.putalpha(layer_pil)
|
||||
|
||||
# 合并到描边层
|
||||
outline_layer = Image.alpha_composite(outline_layer, layer_colored)
|
||||
|
||||
# 创建结果：背景 + 描边 + 人像
|
||||
result = Image.new('RGB', (width, height))
|
||||
result.paste(image, (0, 0))
|
||||
|
||||
# 将描边和人像合成到背景上
|
||||
result_rgba = result.convert('RGBA')
|
||||
result_rgba = Image.alpha_composite(result_rgba, outline_layer)
|
||||
result_rgba = Image.alpha_composite(result_rgba, portrait)
|
||||
image = result_rgba.convert('RGB')
|
||||
|
||||
print("Portrait stroke effect applied successfully", file=sys.stderr)
|
||||
except Exception as e:
|
||||
print(f"Portrait stroke failed: {e}", file=sys.stderr)
|
||||
import traceback
|
||||
traceback.print_exc(file=sys.stderr)
|
||||
else:
|
||||
print("RMBG not available, skipping portrait stroke", file=sys.stderr)
|
||||
|
||||
# 3. 添加蒙版（如果有）
|
||||
mask_config = template.get('mask', {})
|
||||
mask_path = mask_config.get('imagePath', '')
|
||||
if mask_path and os.path.exists(mask_path):
|
||||
try:
|
||||
print(f"Adding mask from: {mask_path}", file=sys.stderr)
|
||||
mask_img = Image.open(mask_path).convert('RGBA')
|
||||
|
||||
# 获取蒙版位置和大小
|
||||
mask_pos = mask_config.get('position', {'x': 10, 'y': 10})
|
||||
mask_size = mask_config.get('size', {'width': 30, 'height': 20})
|
||||
mask_opacity = mask_config.get('opacity', 0.8)
|
||||
|
||||
mask_width = int(mask_size['width'] / 100 * width)
|
||||
mask_height = int(mask_size['height'] / 100 * height)
|
||||
mask_x = int(mask_pos['x'] / 100 * width)
|
||||
mask_y = int(mask_pos['y'] / 100 * height)
|
||||
|
||||
# 调整蒙版大小和透明度
|
||||
mask_img = mask_img.resize((mask_width, mask_height), Image.LANCZOS)
|
||||
if mask_opacity < 1.0:
|
||||
alpha = mask_img.split()[3]
|
||||
alpha = alpha.point(lambda p: int(p * mask_opacity))
|
||||
mask_img.putalpha(alpha)
|
||||
|
||||
# 将蒙版粘贴到图像上
|
||||
image_rgba = image.convert('RGBA')
|
||||
image_rgba.paste(mask_img, (mask_x - mask_width // 2, mask_y - mask_height // 2), mask_img)
|
||||
image = image_rgba.convert('RGB')
|
||||
|
||||
print("Mask applied successfully", file=sys.stderr)
|
||||
except Exception as e:
|
||||
print(f"Failed to add mask: {e}", file=sys.stderr)
|
||||
import traceback
|
||||
traceback.print_exc(file=sys.stderr)
|
||||
|
||||
# 4. 添加标题
|
||||
draw = ImageDraw.Draw(image)
|
||||
|
||||
# 加载字体
|
||||
try:
|
||||
font_main = ImageFont.truetype("NotoSerifCJK-VF", template['titles']['main'].get('fontSize', 60))
|
||||
font_sub = ImageFont.truetype("NotoSerifCJK-VF", template['titles']['sub'].get('fontSize', 40))
|
||||
except:
|
||||
try:
|
||||
font_main = ImageFont.truetype("NotoSerifCJK-VF", template['titles']['main'].get('fontSize', 60))
|
||||
font_sub = ImageFont.truetype("NotoSerifCJK-VF", template['titles']['sub'].get('fontSize', 40))
|
||||
except:
|
||||
font_main = ImageFont.load_default()
|
||||
font_sub = ImageFont.load_default()
|
||||
print("Warning: Failed to load custom font, using default", file=sys.stderr)
|
||||
|
||||
# 绘制主标题（支持自动换行）
|
||||
main_title = template['titles']['main']
|
||||
if main_title.get('text'):
|
||||
main_text = main_title['text']
|
||||
main_color = self._hex_to_rgb(main_title.get('color', '#FFFFFF'))
|
||||
main_pos_x = int(width * main_title['position']['x'] / 100)
|
||||
main_pos_y = int(height * main_title['position']['y'] / 100)
|
||||
max_chars_per_line = main_title.get('maxLength', 4) # 每行最多字符数
|
||||
|
||||
# 将文字分行
|
||||
lines = []
|
||||
for i in range(0, len(main_text), max_chars_per_line):
|
||||
lines.append(main_text[i:i + max_chars_per_line])
|
||||
|
||||
# 计算行高和总高度
|
||||
line_height = int(font_main.size * 1.2) # 行高为字体大小的1.2倍
|
||||
total_height = len(lines) * line_height
|
||||
|
||||
# 计算起始Y坐标，使多行文字垂直居中
|
||||
start_y = main_pos_y - total_height // 2 + line_height // 2
|
||||
|
||||
# 逐行绘制
|
||||
for i, line in enumerate(lines):
|
||||
line_y = start_y + i * line_height
|
||||
|
||||
# 绘制描边
|
||||
if main_title.get('strokeWidth', 0) > 0:
|
||||
stroke_color = self._hex_to_rgb(main_title.get('strokeColor', '#000000'))
|
||||
for offset_x in range(-main_title['strokeWidth'], main_title['strokeWidth'] + 1):
|
||||
for offset_y in range(-main_title['strokeWidth'], main_title['strokeWidth'] + 1):
|
||||
if offset_x == 0 and offset_y == 0:
|
||||
continue
|
||||
draw.text((main_pos_x + offset_x, line_y + offset_y),
|
||||
line, font=font_main, fill=stroke_color, anchor='mm')
|
||||
|
||||
# 绘制主体
|
||||
draw.text((main_pos_x, line_y), line, font=font_main, fill=main_color, anchor='mm')
|
||||
|
||||
print(f"Drew main title: '{main_text}' ({len(lines)} lines) at ({main_pos_x}, {main_pos_y})", file=sys.stderr)
|
||||
|
||||
# 绘制副标题（支持自动换行）
|
||||
sub_title = template['titles']['sub']
|
||||
if sub_title.get('text'):
|
||||
sub_text = sub_title['text']
|
||||
sub_color = self._hex_to_rgb(sub_title.get('color', '#FFFFFF'))
|
||||
sub_pos_x = int(width * sub_title['position']['x'] / 100)
|
||||
sub_pos_y = int(height * sub_title['position']['y'] / 100)
|
||||
max_chars_per_line = sub_title.get('maxLength', 3) # 每行最多字符数
|
||||
|
||||
# 将文字分行
|
||||
lines = []
|
||||
for i in range(0, len(sub_text), max_chars_per_line):
|
||||
lines.append(sub_text[i:i + max_chars_per_line])
|
||||
|
||||
# 计算行高和总高度
|
||||
line_height = int(font_sub.size * 1.2) # 行高为字体大小的1.2倍
|
||||
total_height = len(lines) * line_height
|
||||
|
||||
# 计算起始Y坐标，使多行文字垂直居中
|
||||
start_y = sub_pos_y - total_height // 2 + line_height // 2
|
||||
|
||||
# 逐行绘制
|
||||
for i, line in enumerate(lines):
|
||||
line_y = start_y + i * line_height
|
||||
|
||||
# 绘制描边
|
||||
if sub_title.get('strokeWidth', 0) > 0:
|
||||
stroke_color = self._hex_to_rgb(sub_title.get('strokeColor', '#000000'))
|
||||
for offset_x in range(-sub_title['strokeWidth'], sub_title['strokeWidth'] + 1):
|
||||
for offset_y in range(-sub_title['strokeWidth'], sub_title['strokeWidth'] + 1):
|
||||
if offset_x == 0 and offset_y == 0:
|
||||
continue
|
||||
draw.text((sub_pos_x + offset_x, line_y + offset_y),
|
||||
line, font=font_sub, fill=stroke_color, anchor='mm')
|
||||
|
||||
# 绘制主体
|
||||
draw.text((sub_pos_x, line_y), line, font=font_sub, fill=sub_color, anchor='mm')
|
||||
|
||||
print(f"Drew sub title: '{sub_text}' ({len(lines)} lines) at ({sub_pos_x}, {sub_pos_y})", file=sys.stderr)
|
||||
|
||||
# 保存封面
|
||||
cover_path = os.path.join(self.output_dir, 'cover_custom_template.png')
|
||||
image.save(cover_path, quality=95)
|
||||
|
||||
print(f"Custom template cover generated successfully: {cover_path}", file=sys.stderr)
|
||||
return cover_path
|
||||
|
||||
except Exception as e:
|
||||
print(f"Error generating custom template cover: {e}", file=sys.stderr)
|
||||
import traceback
|
||||
traceback.print_exc(file=sys.stderr)
|
||||
raise Exception(f"Custom template cover generation failed: {str(e)}")
|
||||
@@ -0,0 +1,35 @@
|
||||
"""
|
||||
查找指定字体文件路径
|
||||
"""
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
# 添加模块路径
|
||||
sys.path.insert(0, str(Path(__file__).parent))
|
||||
|
||||
from modules.font_manager import get_font_manager
|
||||
|
||||
def main():
|
||||
if len(sys.argv) < 2:
|
||||
print("Usage: python find_font.py <font_family> [font_weight]", file=sys.stderr)
|
||||
sys.exit(1)
|
||||
|
||||
font_family = sys.argv[1]
|
||||
font_weight = int(sys.argv[2]) if len(sys.argv) > 2 else 400
|
||||
|
||||
try:
|
||||
font_manager = get_font_manager()
|
||||
font_path = font_manager.find_font(font_family, font_weight)
|
||||
|
||||
if font_path:
|
||||
print(font_path)
|
||||
else:
|
||||
print("") # 返回空字符串表示未找到
|
||||
|
||||
sys.exit(0)
|
||||
except Exception as e:
|
||||
print(f"Error: {str(e)}", file=sys.stderr)
|
||||
sys.exit(1)
|
||||
|
||||
if __name__ == '__main__':
|
||||
main()
|
||||
@@ -0,0 +1,81 @@
|
||||
#!/usr/bin/env python3
|
||||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
生成封面模板预览图
|
||||
Generate Cover Template Previews
|
||||
"""
|
||||
|
||||
import os
|
||||
import sys
|
||||
from subtitle_cover_generator_simple import SubtitleCoverGenerator
|
||||
|
||||
# 定义封面样式
|
||||
COVER_STYLES = [
|
||||
'default',
|
||||
'blur-bg',
|
||||
'outline',
|
||||
'gradient',
|
||||
'split'
|
||||
]
|
||||
|
||||
def generate_all_cover_previews():
|
||||
"""生成所有封面模板的预览图"""
|
||||
|
||||
# 使用项目中的demo视频
|
||||
video_path = r"C:\aigcpanel-main\demo\demo-1.mp4"
|
||||
|
||||
if not os.path.exists(video_path):
|
||||
print(f"Error: Video file not found: {video_path}", file=sys.stderr)
|
||||
return False
|
||||
|
||||
# 创建预览图存储目录
|
||||
preview_dir = r"C:\aigcpanel-main\public\cover-previews"
|
||||
os.makedirs(preview_dir, exist_ok=True)
|
||||
|
||||
print(f"Generating cover previews from: {video_path}", file=sys.stderr)
|
||||
print(f"Output directory: {preview_dir}", file=sys.stderr)
|
||||
|
||||
# 为每个样式生成封面预览
|
||||
for style in COVER_STYLES:
|
||||
try:
|
||||
print(f"\nGenerating preview for style: {style}", file=sys.stderr)
|
||||
|
||||
# 创建配置
|
||||
config = {
|
||||
'videoPath': video_path,
|
||||
'scriptText': '欢迎来到我们的精彩世界！这里有无限的可能性和创意。让我们一起探索未知的领域，发现更多的美好事物。', # 使用更长的演示文本
|
||||
'subtitleStyle': 'default',
|
||||
'coverStyle': style,
|
||||
'outputDir': preview_dir,
|
||||
'coverOnly': True # 仅生成封面
|
||||
}
|
||||
|
||||
# 创建生成器并生成封面
|
||||
generator = SubtitleCoverGenerator(config)
|
||||
cover_path = generator.generate_cover()
|
||||
|
||||
# 重命名为样式名称
|
||||
style_preview_path = os.path.join(preview_dir, f'cover-preview-{style}.png')
|
||||
if os.path.exists(cover_path):
|
||||
# 如果目标文件已存在，先删除
|
||||
if os.path.exists(style_preview_path):
|
||||
os.remove(style_preview_path)
|
||||
os.rename(cover_path, style_preview_path)
|
||||
print(f"✓ Generated: {style_preview_path}", file=sys.stderr)
|
||||
else:
|
||||
print(f"✗ Failed to generate: {style}", file=sys.stderr)
|
||||
|
||||
except Exception as e:
|
||||
print(f"✗ Error generating {style}: {e}", file=sys.stderr)
|
||||
continue
|
||||
|
||||
print("\n" + "=" * 60, file=sys.stderr)
|
||||
print("Cover preview generation completed!", file=sys.stderr)
|
||||
print(f"Preview images saved to: {preview_dir}", file=sys.stderr)
|
||||
print("=" * 60, file=sys.stderr)
|
||||
|
||||
return True
|
||||
|
||||
if __name__ == '__main__':
|
||||
success = generate_all_cover_previews()
|
||||
sys.exit(0 if success else 1)
|
||||
@@ -0,0 +1,168 @@
|
||||
#!/usr/bin/env python
|
||||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
为图片生成5种字幕封面风格
|
||||
"""
|
||||
from PIL import Image, ImageDraw, ImageFont
|
||||
import os
|
||||
|
||||
def add_text_with_style(img, title, style_name, font_path):
|
||||
"""在图像上添加指定风格的标题文字"""
|
||||
pil_img = img.copy()
|
||||
draw = ImageDraw.Draw(pil_img)
|
||||
|
||||
width, height = pil_img.size
|
||||
|
||||
# 设置字体大小
|
||||
font_size = int(height * 0.12)
|
||||
try:
|
||||
font = ImageFont.truetype(font_path, font_size)
|
||||
except:
|
||||
font = ImageFont.load_default()
|
||||
|
||||
# 计算标题位置（顶部居中）
|
||||
bbox = draw.textbbox((0, 0), title, font=font)
|
||||
text_width = bbox[2] - bbox[0]
|
||||
text_height = bbox[3] - bbox[1]
|
||||
x = (width - text_width) // 2
|
||||
y = int(height * 0.08) # 距离顶部8%的位置
|
||||
|
||||
if style_name == "default":
|
||||
# 默认风格：黑色背景条 + 白色文字
|
||||
padding = 30
|
||||
draw.rectangle(
|
||||
[x - padding, y - padding, x + text_width + padding, y + text_height + padding],
|
||||
fill=(0, 0, 0, 200)
|
||||
)
|
||||
draw.text((x, y), title, font=font, fill=(255, 255, 255))
|
||||
|
||||
elif style_name == "blur_bg":
|
||||
# 模糊背景风格
|
||||
bg_region = pil_img.crop((x - 50, y - 50, x + text_width + 50, y + text_height + 50))
|
||||
bg_region = bg_region.filter(Image.BLUR)
|
||||
pil_img.paste(bg_region, (x - 50, y - 50))
|
||||
draw = ImageDraw.Draw(pil_img)
|
||||
|
||||
# 添加半透明黑色遮罩
|
||||
overlay = Image.new('RGBA', pil_img.size, (0, 0, 0, 0))
|
||||
overlay_draw = ImageDraw.Draw(overlay)
|
||||
overlay_draw.rectangle([x - 50, y - 50, x + text_width + 50, y + text_height + 50], fill=(0, 0, 0, 150))
|
||||
pil_img = Image.alpha_composite(pil_img.convert('RGBA'), overlay).convert('RGB')
|
||||
draw = ImageDraw.Draw(pil_img)
|
||||
|
||||
draw.text((x, y), title, font=font, fill=(255, 255, 255))
|
||||
|
||||
elif style_name == "outline":
|
||||
# 描边风格：白色文字 + 黑色描边
|
||||
outline_width = 4
|
||||
for adj_x in range(-outline_width, outline_width + 1):
|
||||
for adj_y in range(-outline_width, outline_width + 1):
|
||||
draw.text((x + adj_x, y + adj_y), title, font=font, fill=(0, 0, 0))
|
||||
draw.text((x, y), title, font=font, fill=(255, 255, 255))
|
||||
|
||||
elif style_name == "gradient":
|
||||
# 渐变背景风格
|
||||
gradient_height = text_height + 100
|
||||
gradient = Image.new('RGBA', (width, gradient_height), (0, 0, 0, 0))
|
||||
gradient_draw = ImageDraw.Draw(gradient)
|
||||
for i in range(gradient_height):
|
||||
alpha = int(250 * (i / gradient_height))
|
||||
gradient_draw.rectangle([0, i, width, i + 1], fill=(0, 0, 0, alpha))
|
||||
|
||||
pil_img_rgba = pil_img.convert('RGBA')
|
||||
pil_img_rgba.paste(gradient, (0, y - 50), gradient)
|
||||
pil_img = pil_img_rgba.convert('RGB')
|
||||
draw = ImageDraw.Draw(pil_img)
|
||||
|
||||
draw.text((x, y), title, font=font, fill=(255, 255, 255))
|
||||
|
||||
elif style_name == "split":
|
||||
# 分栏风格：左侧标题 + 右侧装饰
|
||||
padding = 40
|
||||
|
||||
# 左侧标题
|
||||
left_x = padding
|
||||
draw.rectangle(
|
||||
[left_x - 20, y - 20, left_x + text_width + 20, y + text_height + 20],
|
||||
fill=(0, 0, 0, 180)
|
||||
)
|
||||
draw.text((left_x, y), title, font=font, fill=(255, 255, 255))
|
||||
|
||||
# 右侧装饰元素
|
||||
right_text = "✨ 精彩"
|
||||
right_bbox = draw.textbbox((0, 0), right_text, font=font)
|
||||
right_width = right_bbox[2] - right_bbox[0]
|
||||
right_x = width - right_width - padding
|
||||
draw.rectangle(
|
||||
[right_x - 25, y - 25, right_x + right_width + 25, y + text_height + 25],
|
||||
fill=(255, 193, 7, 200) # 金黄色背景
|
||||
)
|
||||
draw.text((right_x, y), right_text, font=font, fill=(0, 0, 0))
|
||||
|
||||
return pil_img
|
||||
|
||||
def main():
|
||||
# 输入图片路径
|
||||
image_path = r"C:\aigcpanel-main\微信图片_20251125215233_633_235.jpg"
|
||||
title = "我是个美女啊啊啊"
|
||||
|
||||
# 输出目录
|
||||
output_dir = r"C:\aigcpanel-main\public\cover-previews"
|
||||
os.makedirs(output_dir, exist_ok=True)
|
||||
|
||||
# 字体路径
|
||||
font_path = r"" # 微软雅黑
|
||||
|
||||
# 读取图片
|
||||
print("正在读取图片...")
|
||||
try:
|
||||
img = Image.open(image_path)
|
||||
print(f"图片尺寸：{img.size}")
|
||||
except Exception as e:
|
||||
print(f"❌ 无法读取图片: {e}")
|
||||
return
|
||||
|
||||
# 调整图片大小以便预览
|
||||
width, height = img.size
|
||||
if width > 1280:
|
||||
scale = 1280 / width
|
||||
new_width = 1280
|
||||
new_height = int(height * scale)
|
||||
img = img.resize((new_width, new_height))
|
||||
print(f"调整图片尺寸为：{img.size}")
|
||||
|
||||
# 定义5种风格
|
||||
styles = [
|
||||
("default", "默认风格：黑色半透明背景条"),
|
||||
("blur_bg", "模糊背景风格：背景模糊效果"),
|
||||
("outline", "描边风格：文字带黑色描边"),
|
||||
("gradient", "渐变背景风格：渐变半透明遮罩"),
|
||||
("split", "分栏风格：标题+装饰元素")
|
||||
]
|
||||
|
||||
# 生成每种风格的封面
|
||||
for style_name, description in styles:
|
||||
print(f"生成 {style_name} 风格封面：{description}")
|
||||
|
||||
# 复制原始图片
|
||||
styled_img = img.copy()
|
||||
|
||||
# 应用风格
|
||||
styled_img = add_text_with_style(styled_img, title, style_name, font_path)
|
||||
|
||||
# 保存封面
|
||||
output_path = os.path.join(output_dir, f"cover-title-{style_name}.png")
|
||||
styled_img.save(output_path)
|
||||
print(f"✅ 已保存：{output_path}")
|
||||
|
||||
print("\n" + "="*60)
|
||||
print("🎉 所有封面生成完成！")
|
||||
print(f"📁 保存位置：{output_dir}")
|
||||
print("="*60)
|
||||
print(f"\n标题：{title}")
|
||||
print("\n封面列表：")
|
||||
for style_name, description in styles:
|
||||
print(f" • cover-title-{style_name}.png - {description}")
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,107 @@
|
||||
"""
|
||||
读取字体文件的内部Family名称（用于ASS字幕）
|
||||
"""
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
# 添加模块路径
|
||||
sys.path.insert(0, str(Path(__file__).parent))
|
||||
|
||||
try:
|
||||
from fontTools.ttLib import TTFont
|
||||
except ImportError:
|
||||
try:
|
||||
import fontTools
|
||||
except ImportError:
|
||||
print("", file=sys.stderr) # 返回空字符串，表示无法读取
|
||||
sys.exit(1)
|
||||
|
||||
def get_font_family_name(font_path: str) -> str:
|
||||
"""
|
||||
获取字体名称。
|
||||
优先使用文件名（不含扩展名），这样可以保证与ziti目录中的文件名匹配。
|
||||
例如：墨趣古风体.ttf → 返回 "墨趣古风体"
|
||||
"""
|
||||
try:
|
||||
# 优先返回文件名（不含扩展名）
|
||||
# 这样可以确保与ziti目录中的文件名匹配，避免使用字体内部的元数据
|
||||
# 字体文件内部的Family Name有时是拼音（如"MoQuGuFengTi"），导致libass找不到字体
|
||||
from pathlib import Path
|
||||
filename = Path(font_path).stem # 获取文件名，不含扩展名
|
||||
if filename:
|
||||
# 确保返回的是正确编码的字符串（对于中文字体名）
|
||||
return str(filename)
|
||||
|
||||
# 如果文件名为空，尝试读取字体内部的Family名称作为备选
|
||||
font = TTFont(font_path)
|
||||
# 读取name表中的Family名称（通常在第1或第16条记录中）
|
||||
for record in font['name'].names:
|
||||
# nameID 1 = Font Family Name
|
||||
if record.nameID == 1:
|
||||
# 尝试解码，优先使用Unicode编码
|
||||
try:
|
||||
if hasattr(record, 'string'):
|
||||
if isinstance(record.string, bytes):
|
||||
return record.string.decode('utf-16-be') if record.isUnicode() else record.string.decode('latin-1')
|
||||
return str(record.string)
|
||||
except:
|
||||
pass
|
||||
# 如果找不到nameID 1，尝试其他nameID
|
||||
for record in font['name'].names:
|
||||
if record.nameID == 16: # Typographic Family Name
|
||||
try:
|
||||
if hasattr(record, 'string'):
|
||||
if isinstance(record.string, bytes):
|
||||
return record.string.decode('utf-16-be') if record.isUnicode() else record.string.decode('latin-1')
|
||||
return str(record.string)
|
||||
except:
|
||||
pass
|
||||
font.close()
|
||||
except Exception as e:
|
||||
pass
|
||||
|
||||
return ""
|
||||
|
||||
def main():
|
||||
if len(sys.argv) < 2:
|
||||
print("", file=sys.stderr)
|
||||
sys.exit(1)
|
||||
|
||||
font_path = sys.argv[1]
|
||||
|
||||
if not Path(font_path).exists():
|
||||
print("", file=sys.stderr)
|
||||
sys.exit(1)
|
||||
|
||||
try:
|
||||
family_name = get_font_family_name(font_path)
|
||||
if family_name:
|
||||
print(family_name)
|
||||
else:
|
||||
print("") # 返回空字符串表示无法读取
|
||||
sys.exit(0)
|
||||
except Exception as e:
|
||||
print("", file=sys.stderr)
|
||||
sys.exit(1)
|
||||
|
||||
if __name__ == '__main__':
|
||||
main()
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
@@ -0,0 +1,26 @@
|
||||
"""
|
||||
获取所有可用字体列表
|
||||
"""
|
||||
import json
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
# 添加模块路径
|
||||
sys.path.insert(0, str(Path(__file__).parent))
|
||||
|
||||
from modules.font_manager import get_font_manager
|
||||
|
||||
def main():
|
||||
try:
|
||||
font_manager = get_font_manager()
|
||||
fonts = font_manager.get_all_available_fonts()
|
||||
|
||||
# 输出JSON格式
|
||||
print(json.dumps(fonts, ensure_ascii=False))
|
||||
sys.exit(0)
|
||||
except Exception as e:
|
||||
print(f"Error: {str(e)}", file=sys.stderr)
|
||||
sys.exit(1)
|
||||
|
||||
if __name__ == '__main__':
|
||||
main()
|
||||
@@ -0,0 +1,100 @@
|
||||
#!/usr/bin/env python
|
||||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
初始化预设封面模板到数据库
|
||||
"""
|
||||
|
||||
import json
|
||||
import sqlite3
|
||||
import os
|
||||
import sys
|
||||
from datetime import datetime
|
||||
|
||||
def get_db_path():
|
||||
"""获取数据库路径"""
|
||||
# 假设数据库在项目根目录的data文件夹
|
||||
project_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
||||
db_path = os.path.join(project_root, 'data', 'aigcpanel.db')
|
||||
return db_path
|
||||
|
||||
def init_templates():
|
||||
"""初始化预设模板"""
|
||||
# 读取模板配置
|
||||
script_dir = os.path.dirname(os.path.abspath(__file__))
|
||||
template_file = os.path.join(script_dir, 'preset_templates.json')
|
||||
|
||||
print(f"Reading templates from: {template_file}")
|
||||
|
||||
with open(template_file, 'r', encoding='utf-8') as f:
|
||||
data = json.load(f)
|
||||
|
||||
templates = data['templates']
|
||||
print(f"Found {len(templates)} preset templates")
|
||||
|
||||
# 连接数据库
|
||||
db_path = get_db_path()
|
||||
print(f"Connecting to database: {db_path}")
|
||||
|
||||
if not os.path.exists(db_path):
|
||||
print(f"Error: Database not found at {db_path}")
|
||||
print("Please ensure the application has been run at least once to create the database.")
|
||||
return False
|
||||
|
||||
conn = sqlite3.connect(db_path)
|
||||
cursor = conn.cursor()
|
||||
|
||||
# 检查表是否存在
|
||||
cursor.execute("""
|
||||
SELECT name FROM sqlite_master
|
||||
WHERE type='table' AND name='cover_templates'
|
||||
""")
|
||||
|
||||
if not cursor.fetchone():
|
||||
print("Error: cover_templates table does not exist")
|
||||
print("Please run the application first to create the database schema.")
|
||||
conn.close()
|
||||
return False
|
||||
|
||||
# 清空现有预设模板（可选）
|
||||
print("Clearing existing preset templates...")
|
||||
cursor.execute("DELETE FROM cover_templates WHERE id LIKE 'preset-%'")
|
||||
|
||||
# 插入预设模板
|
||||
now = int(datetime.now().timestamp() * 1000)
|
||||
|
||||
for idx, template in enumerate(templates, 1):
|
||||
template_id = f"preset-{idx:02d}"
|
||||
name = template['name']
|
||||
config_json = json.dumps(template['config'], ensure_ascii=False)
|
||||
|
||||
print(f"Inserting template: {name} (ID: {template_id})")
|
||||
|
||||
cursor.execute("""
|
||||
INSERT INTO cover_templates
|
||||
(id, name, config, thumbnail_path, created_at, updated_at)
|
||||
VALUES (?, ?, ?, NULL, ?, ?)
|
||||
""", (template_id, name, config_json, now, now))
|
||||
|
||||
conn.commit()
|
||||
conn.close()
|
||||
|
||||
print(f"\n✅ Successfully initialized {len(templates)} preset templates!")
|
||||
return True
|
||||
|
||||
def main():
|
||||
print("="*60)
|
||||
print("Preset Cover Templates Initialization")
|
||||
print("="*60)
|
||||
print()
|
||||
|
||||
try:
|
||||
success = init_templates()
|
||||
sys.exit(0 if success else 1)
|
||||
except Exception as e:
|
||||
print(f"\n❌ Error: {e}")
|
||||
import traceback
|
||||
traceback.print_exc()
|
||||
sys.exit(1)
|
||||
|
||||
if __name__ == '__main__':
|
||||
main()
|
||||
@@ -0,0 +1,6 @@
|
||||
"""
|
||||
modules 包初始化文件
|
||||
"""
|
||||
|
||||
# 这个文件是必需的，让 Python 将 modules 目录识别为一个包
|
||||
# 从而支持相对导入（如 from .utils import ...）
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,117 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
封面模板定义文件
|
||||
Cover Template Definitions
|
||||
"""
|
||||
|
||||
from typing import Dict, Any, Tuple
|
||||
|
||||
# 封面模板配置
|
||||
COVER_TEMPLATES: Dict[str, Dict[str, Any]] = {
|
||||
'professional': {
|
||||
'name': '专业商务风',
|
||||
'description': '金色描边 + 中度虚化 + 人物清晰',
|
||||
'blur_radius': 25, # 背景模糊半径
|
||||
'outline_color': (255, 215, 0), # 金色 (R, G, B)
|
||||
'outline_width': 8, # 描边宽度
|
||||
'brightness': 0.9, # 亮度调整 (0.0-1.0)
|
||||
'contrast': 1.0, # 对比度
|
||||
'preview_path': '/assets/templates/professional.jpg'
|
||||
},
|
||||
'vibrant': {
|
||||
'name': '活力青春风',
|
||||
'description': '粉色描边 + 轻度虚化 + 明亮色调',
|
||||
'blur_radius': 20,
|
||||
'outline_color': (255, 105, 180), # 粉色 (R, G, B)
|
||||
'outline_width': 6,
|
||||
'brightness': 1.1,
|
||||
'contrast': 1.05,
|
||||
'preview_path': '/assets/templates/vibrant.jpg'
|
||||
},
|
||||
'elegant': {
|
||||
'name': '优雅高级风',
|
||||
'description': '银色描边 + 重度虚化 + 柔和光效',
|
||||
'blur_radius': 30,
|
||||
'outline_color': (192, 192, 192), # 银色 (R, G, B)
|
||||
'outline_width': 10,
|
||||
'brightness': 0.85,
|
||||
'contrast': 0.95,
|
||||
'preview_path': '/assets/templates/elegant.jpg'
|
||||
},
|
||||
'classic': {
|
||||
'name': '经典商务风',
|
||||
'description': '白色描边 + 轻度虚化 + 简洁风格',
|
||||
'blur_radius': 15,
|
||||
'outline_color': (255, 255, 255), # 白色 (R, G, B)
|
||||
'outline_width': 5,
|
||||
'brightness': 0.95,
|
||||
'contrast': 1.0,
|
||||
'preview_path': '/assets/templates/classic.jpg'
|
||||
},
|
||||
'dramatic': {
|
||||
'name': '戏剧艺术风',
|
||||
'description': '红色描边 + 重度虚化 + 强烈对比',
|
||||
'blur_radius': 35,
|
||||
'outline_color': (220, 20, 60), # 红色 (R, G, B)
|
||||
'outline_width': 12,
|
||||
'brightness': 0.8,
|
||||
'contrast': 1.2,
|
||||
'preview_path': '/assets/templates/dramatic.jpg'
|
||||
}
|
||||
}
|
||||
|
||||
def get_template(template_id: str) -> Dict[str, Any]:
|
||||
"""
|
||||
获取模板配置
|
||||
|
||||
Args:
|
||||
template_id: 模板ID
|
||||
|
||||
Returns:
|
||||
模板配置字典
|
||||
"""
|
||||
return COVER_TEMPLATES.get(template_id, COVER_TEMPLATES['professional'])
|
||||
|
||||
def get_all_templates() -> Dict[str, Dict[str, Any]]:
|
||||
"""
|
||||
获取所有模板配置
|
||||
|
||||
Returns:
|
||||
所有模板配置字典
|
||||
"""
|
||||
return COVER_TEMPLATES.copy()
|
||||
|
||||
def get_template_names() -> Dict[str, str]:
|
||||
"""
|
||||
获取模板ID和名称映射
|
||||
|
||||
Returns:
|
||||
{template_id: template_name}
|
||||
"""
|
||||
return {k: v['name'] for k, v in COVER_TEMPLATES.items()}
|
||||
|
||||
def get_template_descriptions() -> Dict[str, str]:
|
||||
"""
|
||||
获取模板ID和描述映射
|
||||
|
||||
Returns:
|
||||
{template_id: template_description}
|
||||
"""
|
||||
return {k: v['description'] for k, v in COVER_TEMPLATES.items()}
|
||||
|
||||
def get_template_preview_paths() -> Dict[str, str]:
|
||||
"""
|
||||
获取模板ID和预览图路径映射
|
||||
|
||||
Returns:
|
||||
{template_id: preview_path}
|
||||
"""
|
||||
return {k: v['preview_path'] for k, v in COVER_TEMPLATES.items()}
|
||||
|
||||
# 模板效果参数范围定义（用于UI验证）
|
||||
TEMPLATE_PARAM_RANGES = {
|
||||
'blur_radius': {'min': 0, 'max': 50, 'default': 25},
|
||||
'outline_width': {'min': 0, 'max': 20, 'default': 8},
|
||||
'brightness': {'min': 0.5, 'max': 1.5, 'default': 1.0},
|
||||
'contrast': {'min': 0.5, 'max': 1.5, 'default': 1.0}
|
||||
}
|
||||
@@ -0,0 +1,396 @@
|
||||
"""
|
||||
字体管理模块
|
||||
负责查找和管理字体文件，支持预置字体、系统字体、ziti目录字体和Google Fonts
|
||||
"""
|
||||
import os
|
||||
import sys
|
||||
import platform
|
||||
from pathlib import Path
|
||||
from typing import Dict, List, Optional, Tuple
|
||||
import json
|
||||
|
||||
# 尝试导入loguru，如果不可用则使用print作为替代
|
||||
try:
|
||||
from loguru import logger
|
||||
except ImportError:
|
||||
# Fallback logger that uses print
|
||||
class logger:
|
||||
@staticmethod
|
||||
def info(msg, *args, **kwargs):
|
||||
print(f"[INFO] {msg}", file=sys.stderr)
|
||||
@staticmethod
|
||||
def warning(msg, *args, **kwargs):
|
||||
print(f"[WARN] {msg}", file=sys.stderr)
|
||||
@staticmethod
|
||||
def error(msg, *args, **kwargs):
|
||||
print(f"[ERROR] {msg}", file=sys.stderr)
|
||||
@staticmethod
|
||||
def debug(msg, *args, **kwargs):
|
||||
print(f"[DEBUG] {msg}", file=sys.stderr)
|
||||
|
||||
|
||||
class FontManager:
|
||||
"""字体管理器"""
|
||||
|
||||
def __init__(self):
|
||||
# 计算项目根目录（从 python/modules 向上两级）
|
||||
# __file__ 应该是 .../python/modules/font_manager.py
|
||||
current_file = Path(__file__).resolve()
|
||||
python_dir = current_file.parent.parent # python 目录
|
||||
project_root = python_dir.parent # 项目根目录
|
||||
|
||||
self.bundled_fonts_dir = Path("fonts/bundled")
|
||||
self.fonts_metadata_file = self.bundled_fonts_dir / "fonts_metadata.json"
|
||||
|
||||
# 🔧 修复ziti字体路径：支持ASAR打包环境
|
||||
# 检查是否在打包环境中运行
|
||||
app_root = os.environ.get('APP_ROOT', None)
|
||||
if app_root:
|
||||
ziti_path_bundle = Path(app_root) / "resources-bundles" / "ziti"
|
||||
if ziti_path_bundle.exists():
|
||||
self.ziti_fonts_dir = ziti_path_bundle
|
||||
print(f"[font_manager] ziti bundle dir: {ziti_path_bundle}", file=sys.stderr)
|
||||
else:
|
||||
ziti_path = Path(app_root) / "extra" / "common" / "fonts" / "ziti"
|
||||
if ziti_path.exists():
|
||||
self.ziti_fonts_dir = ziti_path
|
||||
print(f"[font_manager] ziti dir: {ziti_path}", file=sys.stderr)
|
||||
else:
|
||||
ziti_path_fallback1 = Path(app_root) / "app.asar.unpacked" / "ziti"
|
||||
if ziti_path_fallback1.exists():
|
||||
self.ziti_fonts_dir = ziti_path_fallback1
|
||||
print(f"[font_manager] ziti fallback dir: {ziti_path_fallback1}", file=sys.stderr)
|
||||
else:
|
||||
ziti_path_fallback2 = Path(app_root) / "ziti"
|
||||
self.ziti_fonts_dir = ziti_path_fallback2
|
||||
print(f"[font_manager] ziti fallback path: {ziti_path_fallback2} (exists: {ziti_path_fallback2.exists()})", file=sys.stderr)
|
||||
else:
|
||||
self.ziti_fonts_dir = project_root / "ziti"
|
||||
print(f"[font_manager] dev ziti path: {self.ziti_fonts_dir}", file=sys.stderr)
|
||||
self._bundled_fonts_cache: Optional[Dict] = None
|
||||
self._system_fonts_cache: Optional[List[Dict]] = None
|
||||
self._ziti_fonts_cache: Optional[List[Dict]] = None # 新增：ziti字体缓存
|
||||
|
||||
# 诊断日志
|
||||
logger.debug(f"FontManager 初始化: current_file={current_file}, python_dir={python_dir}, project_root={project_root}")
|
||||
logger.debug(f"ziti_fonts_dir={self.ziti_fonts_dir.absolute()}, exists={self.ziti_fonts_dir.exists()}")
|
||||
|
||||
def get_bundled_fonts(self) -> List[Dict]:
|
||||
"""获取预置字体列表"""
|
||||
if self._bundled_fonts_cache is not None:
|
||||
return self._bundled_fonts_cache.get("fonts", [])
|
||||
|
||||
if not self.fonts_metadata_file.exists():
|
||||
logger.warning(f"字体元数据文件不存在: {self.fonts_metadata_file}")
|
||||
return []
|
||||
|
||||
try:
|
||||
with open(self.fonts_metadata_file, "r", encoding="utf-8") as f:
|
||||
metadata = json.load(f)
|
||||
self._bundled_fonts_cache = metadata
|
||||
return metadata.get("fonts", [])
|
||||
except Exception as e:
|
||||
logger.error(f"读取字体元数据失败: {e}")
|
||||
return []
|
||||
|
||||
def scan_ziti_fonts(self) -> List[Dict]:
|
||||
"""扫描ziti目录中的字体文件"""
|
||||
if self._ziti_fonts_cache is not None:
|
||||
return self._ziti_fonts_cache
|
||||
|
||||
ziti_fonts = []
|
||||
|
||||
logger.info(f"正在扫描 ziti 目录: {self.ziti_fonts_dir.absolute()}")
|
||||
logger.info(f" 当前工作目录: {os.getcwd()}")
|
||||
logger.info(f" APP_ROOT环境变量: {os.environ.get('APP_ROOT', 'None')}")
|
||||
|
||||
if not self.ziti_fonts_dir.exists():
|
||||
logger.warning(f"❌ ziti目录不存在: {self.ziti_fonts_dir.absolute()}")
|
||||
logger.warning(f" 请检查路径配置和资源文件是否正确打包")
|
||||
return []
|
||||
else:
|
||||
logger.info(f"✅ ziti目录存在，开始扫描字体文件...")
|
||||
|
||||
try:
|
||||
# 支持的字体格式
|
||||
font_extensions = ['.ttf', '.otf', '.ttc', '.woff', '.woff2']
|
||||
|
||||
# 扫描ziti目录中的所有字体文件
|
||||
for font_file in self.ziti_fonts_dir.iterdir():
|
||||
if font_file.is_file() and font_file.suffix.lower() in font_extensions:
|
||||
# 从文件名提取字体信息
|
||||
font_name = font_file.stem
|
||||
|
||||
# 尝试解析字体名称和变体
|
||||
display_name = font_name
|
||||
family = font_name
|
||||
weight = 400
|
||||
|
||||
# 检测常见的字体变体
|
||||
name_lower = font_name.lower()
|
||||
if 'bold' in name_lower:
|
||||
weight = 700
|
||||
elif 'semibold' in name_lower:
|
||||
weight = 600
|
||||
elif 'medium' in name_lower:
|
||||
weight = 500
|
||||
elif 'light' in name_lower:
|
||||
weight = 300
|
||||
elif 'thin' in name_lower:
|
||||
weight = 100
|
||||
elif 'black' in name_lower or 'extrabold' in name_lower:
|
||||
weight = 900
|
||||
|
||||
# 清理family名称（移除变体后缀）
|
||||
for variant in ['-Bold', '-SemiBold', '-Medium', '-Light', '-Thin', '-Black', '-ExtraBold', '-Regular']:
|
||||
if family.endswith(variant):
|
||||
family = family[:-len(variant)]
|
||||
break
|
||||
|
||||
ziti_fonts.append({
|
||||
'family': family,
|
||||
'display_name': display_name,
|
||||
'path': str(font_file.absolute()),
|
||||
'source': 'ziti',
|
||||
'weight': weight,
|
||||
'category': 'sans-serif', # 默认分类
|
||||
'languages': ['zh-CN', 'ja', 'en'] # 假设支持中日英
|
||||
})
|
||||
|
||||
logger.info(f"✅ 从ziti目录扫描到 {len(ziti_fonts)} 个字体")
|
||||
if len(ziti_fonts) > 0:
|
||||
logger.debug(f" 扫描到的字体: {', '.join([f['family'] for f in ziti_fonts[:5]])}" +
|
||||
(f" 等({len(ziti_fonts)}个)" if len(ziti_fonts) > 5 else ""))
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"❌ 扫描ziti目录字体失败: {e}")
|
||||
import traceback
|
||||
logger.error(f" 追踪: {traceback.format_exc()}")
|
||||
|
||||
# 去重（按family和weight）
|
||||
seen = set()
|
||||
unique_fonts = []
|
||||
for font in ziti_fonts:
|
||||
key = (font['family'], font.get('weight', 400))
|
||||
if key not in seen:
|
||||
seen.add(key)
|
||||
unique_fonts.append(font)
|
||||
|
||||
self._ziti_fonts_cache = unique_fonts
|
||||
return unique_fonts
|
||||
|
||||
def scan_system_fonts(self) -> List[Dict]:
|
||||
"""Do not scan OS fonts. Runtime must use bundled font package only."""
|
||||
self._system_fonts_cache = []
|
||||
return []
|
||||
def find_font(self, font_family: str, font_weight: int = 400) -> Optional[str]:
|
||||
"""Find a font file from bundled metadata or ziti only."""
|
||||
bundled_font = self._find_bundled_font(font_family, font_weight)
|
||||
if bundled_font:
|
||||
return bundled_font
|
||||
|
||||
ziti_font = self._find_ziti_font(font_family, font_weight)
|
||||
if ziti_font:
|
||||
return ziti_font
|
||||
|
||||
return None
|
||||
def _find_bundled_font(self, font_family: str, font_weight: int) -> Optional[str]:
|
||||
"""查找预置字体"""
|
||||
if not self.bundled_fonts_dir.exists():
|
||||
return None
|
||||
|
||||
bundled_fonts = self.get_bundled_fonts()
|
||||
|
||||
# 查找匹配的字体
|
||||
for font_info in bundled_fonts:
|
||||
if font_info.get('family') == font_family:
|
||||
font_dir = self.bundled_fonts_dir / font_info.get('path', '')
|
||||
|
||||
# 根据font_weight选择变体
|
||||
variant = self._get_variant_for_weight(font_weight)
|
||||
|
||||
# 查找字体文件（支持TTF和OTF格式）
|
||||
for ext in [".ttf", ".otf"]:
|
||||
variant_file = font_dir / f"{variant}{ext}"
|
||||
if variant_file.exists():
|
||||
return str(variant_file)
|
||||
|
||||
# 如果找不到指定变体，尝试查找Regular
|
||||
for ext in [".ttf", ".otf"]:
|
||||
regular_file = font_dir / f"Regular{ext}"
|
||||
if regular_file.exists():
|
||||
return str(regular_file)
|
||||
|
||||
return None
|
||||
|
||||
def _find_ziti_font(self, font_family: str, font_weight: int) -> Optional[str]:
|
||||
"""查找ziti目录字体（支持模糊匹配和名称映射）"""
|
||||
if not font_family:
|
||||
return None
|
||||
|
||||
ziti_fonts = self.scan_ziti_fonts()
|
||||
|
||||
if not ziti_fonts:
|
||||
logger.debug(f"ziti目录中没有字体文件，无法查找: {font_family}")
|
||||
return None
|
||||
|
||||
# 字体名称映射表（与前端保持一致）
|
||||
# 映射：显示名称/别名 -> 文件名（不包含扩展名）
|
||||
font_name_map = {
|
||||
'Dymon手写体': 'Dymon-ShouXieTi',
|
||||
'Dymon-ShouXieTi': 'Dymon-ShouXieTi',
|
||||
'猫啃杂糅体': 'MaokenAssortedSans',
|
||||
'MaokenAssortedSans': 'MaokenAssortedSans',
|
||||
'猫啃杂糅体 Lite': 'MaokenAssortedSans-Lite',
|
||||
'MaokenAssortedSans-Lite': 'MaokenAssortedSans-Lite',
|
||||
'Murecho 黑体': 'Murecho-Black',
|
||||
'Murecho-Black': 'Murecho-Black',
|
||||
'Murecho 粗体': 'Murecho-Bold',
|
||||
'Murecho-Bold': 'Murecho-Bold',
|
||||
'墨趣古风体': '墨趣古风体',
|
||||
'平方张亚玲黑方体': '平方张亚玲黑方体',
|
||||
'胡晓波骚包体': '胡晓波骚包体2.0',
|
||||
'胡晓波骚包体2.0': '胡晓波骚包体2.0',
|
||||
}
|
||||
|
||||
# 标准化字体名称（移除空格、统一大小写）
|
||||
def normalize_name(name: str) -> str:
|
||||
name = name.lower()
|
||||
for ext in ['.ttf', '.otf', '.ttc', '.woff', '.woff2']:
|
||||
while name.endswith(ext):
|
||||
name = name[:-len(ext)]
|
||||
return name.replace(' ', '').replace('-', '').replace('_', '').lower()
|
||||
|
||||
# 尝试通过映射表转换字体名称
|
||||
mapped_font_family = font_name_map.get(font_family, font_family)
|
||||
normalized_target = normalize_name(mapped_font_family)
|
||||
|
||||
logger.debug(f"查找ziti字体: '{font_family}' -> 映射: '{mapped_font_family}' -> 标准化: '{normalized_target}', 权重: {font_weight}")
|
||||
|
||||
# 查找匹配的字体
|
||||
best_match = None
|
||||
min_weight_diff = float('inf')
|
||||
|
||||
for font_info in ziti_fonts:
|
||||
family = font_info.get('family', '')
|
||||
display_name = font_info.get('display_name', '')
|
||||
path = font_info.get('path', '')
|
||||
|
||||
# 获取文件名（不含扩展名）用于匹配
|
||||
file_stem = None
|
||||
if path:
|
||||
from pathlib import Path
|
||||
file_stem = Path(path).stem
|
||||
|
||||
# 多种匹配方式：精确匹配、模糊匹配、文件名匹配
|
||||
is_match = False
|
||||
match_type = None
|
||||
|
||||
# 1. 精确匹配
|
||||
if family == font_family or family == mapped_font_family:
|
||||
is_match = True
|
||||
match_type = f"精确匹配(family={family})"
|
||||
elif display_name == font_family or display_name == mapped_font_family:
|
||||
is_match = True
|
||||
match_type = f"精确匹配(display_name={display_name})"
|
||||
# 2. 模糊匹配（忽略大小写和空格）
|
||||
elif normalize_name(family) == normalized_target:
|
||||
is_match = True
|
||||
match_type = f"模糊匹配(family={family})"
|
||||
elif normalize_name(display_name) == normalized_target:
|
||||
is_match = True
|
||||
match_type = f"模糊匹配(display_name={display_name})"
|
||||
# 3. 文件名匹配
|
||||
elif file_stem and normalize_name(file_stem) == normalized_target:
|
||||
is_match = True
|
||||
match_type = f"文件名匹配(file_stem={file_stem})"
|
||||
|
||||
if is_match:
|
||||
font_weight_info = font_info.get('weight', 400)
|
||||
weight_diff = abs(font_weight_info - font_weight)
|
||||
|
||||
logger.debug(f"找到匹配字体: {match_type}, 路径: {path}, 权重: {font_weight_info}, 权重差: {weight_diff}")
|
||||
|
||||
# 找到权重最接近的字体
|
||||
if weight_diff < min_weight_diff:
|
||||
min_weight_diff = weight_diff
|
||||
best_match = path
|
||||
|
||||
if best_match:
|
||||
logger.info(f"找到ziti字体 '{font_family}': {best_match}")
|
||||
else:
|
||||
logger.warning(f"未找到ziti字体 '{font_family}'，已扫描{len(ziti_fonts)}个字体文件")
|
||||
|
||||
return best_match
|
||||
|
||||
def _find_system_font(self, font_family: str, font_weight: int) -> Optional[str]:
|
||||
"""查找系统字体"""
|
||||
system_fonts = self.scan_system_fonts()
|
||||
|
||||
# 查找匹配的字体
|
||||
for font_info in system_fonts:
|
||||
if font_info.get('family') == font_family:
|
||||
font_weight_info = font_info.get('weight', 400)
|
||||
|
||||
# 如果权重匹配（或接近），返回字体路径
|
||||
if abs(font_weight_info - font_weight) <= 100:
|
||||
return font_info.get('path')
|
||||
|
||||
return None
|
||||
|
||||
def _get_variant_for_weight(self, font_weight: int) -> str:
|
||||
"""根据字体权重获取变体名称"""
|
||||
if font_weight >= 700:
|
||||
return "Bold"
|
||||
elif font_weight >= 600:
|
||||
return "SemiBold"
|
||||
elif font_weight >= 500:
|
||||
return "Medium"
|
||||
else:
|
||||
return "Regular"
|
||||
|
||||
def get_all_available_fonts(self) -> List[Dict]:
|
||||
"""获取所有可用字体（合并预置、ziti和系统字体）"""
|
||||
all_fonts = []
|
||||
|
||||
# 添加预置字体
|
||||
bundled_fonts = self.get_bundled_fonts()
|
||||
for font_info in bundled_fonts:
|
||||
all_fonts.append({
|
||||
'value': font_info.get('family'),
|
||||
'label': font_info.get('display_name', font_info.get('family')),
|
||||
'source': 'bundled',
|
||||
'category': font_info.get('category', 'sans-serif'),
|
||||
'languages': font_info.get('languages', [])
|
||||
})
|
||||
|
||||
# 添加ziti目录字体
|
||||
ziti_fonts = self.scan_ziti_fonts()
|
||||
seen_families = {f.get('family') for f in bundled_fonts}
|
||||
|
||||
for font_info in ziti_fonts:
|
||||
family = font_info.get('family')
|
||||
if family not in seen_families:
|
||||
all_fonts.append({
|
||||
'value': family,
|
||||
'label': font_info.get('display_name', family),
|
||||
'source': 'ziti',
|
||||
'category': font_info.get('category', 'sans-serif'),
|
||||
'languages': font_info.get('languages', []),
|
||||
'path': font_info.get('path')
|
||||
})
|
||||
seen_families.add(family)
|
||||
|
||||
return all_fonts
|
||||
|
||||
|
||||
# 全局字体管理器实例
|
||||
_font_manager: Optional[FontManager] = None
|
||||
|
||||
|
||||
def get_font_manager() -> FontManager:
|
||||
"""获取全局字体管理器实例"""
|
||||
global _font_manager
|
||||
if _font_manager is None:
|
||||
_font_manager = FontManager()
|
||||
return _font_manager
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,312 @@
|
||||
"""
|
||||
工具函数模块
|
||||
提供字幕封面生成过程中需要的通用工具函数
|
||||
"""
|
||||
|
||||
import os
|
||||
import sys
|
||||
from pathlib import Path
|
||||
from typing import Tuple, Optional, List, Dict, Any
|
||||
|
||||
# 尝试导入loguru，如果不可用则使用print作为替代
|
||||
try:
|
||||
from loguru import logger
|
||||
except ImportError:
|
||||
# Fallback logger that uses print
|
||||
class logger:
|
||||
@staticmethod
|
||||
def info(msg, *args, **kwargs):
|
||||
print(f"[INFO] {msg}", file=sys.stderr)
|
||||
@staticmethod
|
||||
def warning(msg, *args, **kwargs):
|
||||
print(f"[WARN] {msg}", file=sys.stderr)
|
||||
@staticmethod
|
||||
def error(msg, *args, **kwargs):
|
||||
print(f"[ERROR] {msg}", file=sys.stderr)
|
||||
@staticmethod
|
||||
def debug(msg, *args, **kwargs):
|
||||
print(f"[DEBUG] {msg}", file=sys.stderr)
|
||||
|
||||
# 尝试导入图像处理库
|
||||
try:
|
||||
import cv2
|
||||
import numpy as np
|
||||
CV2_AVAILABLE = True
|
||||
except ImportError:
|
||||
CV2_AVAILABLE = False
|
||||
logger.warning("opencv-python不可用，图像处理功能受限")
|
||||
|
||||
try:
|
||||
from PIL import Image
|
||||
PIL_AVAILABLE = True
|
||||
except ImportError:
|
||||
PIL_AVAILABLE = False
|
||||
logger.warning("PIL不可用，图像处理功能受限")
|
||||
|
||||
|
||||
def ensure_dir(path: str) -> Path:
|
||||
"""确保目录存在"""
|
||||
path_obj = Path(path)
|
||||
path_obj.mkdir(parents=True, exist_ok=True)
|
||||
return path_obj
|
||||
|
||||
|
||||
def get_video_info(video_path: str) -> Dict[str, Any]:
|
||||
"""获取视频信息"""
|
||||
if not CV2_AVAILABLE:
|
||||
# 如果opencv不可用，返回基本信息
|
||||
try:
|
||||
import os
|
||||
file_size = os.path.getsize(video_path)
|
||||
return {
|
||||
'fps': 30, # 默认值
|
||||
'frame_count': 0,
|
||||
'width': 1920, # 默认值
|
||||
'height': 1080, # 默认值
|
||||
'duration': 0, # 未知
|
||||
'path': video_path,
|
||||
'file_size': file_size,
|
||||
'note': '视频信息不完整，opencv不可用'
|
||||
}
|
||||
except Exception as e:
|
||||
logger.error(f"获取基本视频信息失败: {e}")
|
||||
raise
|
||||
|
||||
try:
|
||||
cap = cv2.VideoCapture(video_path)
|
||||
if not cap.isOpened():
|
||||
raise ValueError(f"无法打开视频文件: {video_path}")
|
||||
|
||||
# 获取视频属性
|
||||
fps = cap.get(cv2.CAP_PROP_FPS)
|
||||
frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
|
||||
width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
|
||||
height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
|
||||
duration = frame_count / fps if fps > 0 else 0
|
||||
|
||||
cap.release()
|
||||
|
||||
return {
|
||||
'fps': fps,
|
||||
'frame_count': frame_count,
|
||||
'width': width,
|
||||
'height': height,
|
||||
'duration': duration,
|
||||
'path': video_path
|
||||
}
|
||||
except Exception as e:
|
||||
logger.error(f"获取视频信息失败: {e}")
|
||||
raise
|
||||
|
||||
|
||||
def extract_frame_at_time(video_path: str, time_seconds: float):
|
||||
"""在指定时间抽取视频帧"""
|
||||
if not CV2_AVAILABLE:
|
||||
logger.warning("opencv不可用，无法抽取视频帧")
|
||||
# 返回一个占位符
|
||||
return None
|
||||
|
||||
try:
|
||||
cap = cv2.VideoCapture(video_path)
|
||||
if not cap.isOpened():
|
||||
raise ValueError(f"无法打开视频文件: {video_path}")
|
||||
|
||||
# 设置帧位置
|
||||
fps = cap.get(cv2.CAP_PROP_FPS)
|
||||
frame_number = int(time_seconds * fps)
|
||||
cap.set(cv2.CAP_PROP_POS_FRAMES, frame_number)
|
||||
|
||||
ret, frame = cap.read()
|
||||
cap.release()
|
||||
|
||||
if not ret:
|
||||
raise ValueError(f"无法读取帧: {frame_number}")
|
||||
|
||||
return frame
|
||||
except Exception as e:
|
||||
logger.error(f"抽帧失败: {e}")
|
||||
raise
|
||||
|
||||
|
||||
def save_image(image, output_path: str, quality: int = 95) -> None:
|
||||
"""保存图像
|
||||
|
||||
如果输出路径是PNG格式，会保留alpha通道（透明背景）
|
||||
如果是JPEG格式，会转换为RGB格式
|
||||
"""
|
||||
try:
|
||||
ensure_dir(os.path.dirname(output_path))
|
||||
|
||||
is_png = output_path.lower().endswith('.png')
|
||||
|
||||
if CV2_AVAILABLE and isinstance(image, np.ndarray):
|
||||
# 使用OpenCV保存
|
||||
if is_png and image.shape[2] == 4:
|
||||
# PNG格式且有alpha通道，保存为BGRA（OpenCV使用BGR格式）
|
||||
# 确保图像是BGRA格式（B, G, R, A）
|
||||
# 如果输入是RGBA（R, G, B, A），需要转换为BGRA
|
||||
if image.dtype != np.uint8:
|
||||
image = image.astype(np.uint8)
|
||||
# OpenCV的imwrite会自动处理BGRA格式
|
||||
success = cv2.imwrite(output_path, image, [cv2.IMWRITE_PNG_COMPRESSION, 3])
|
||||
elif is_png:
|
||||
# PNG格式但没有alpha通道，转换为RGB
|
||||
if image.shape[2] == 3:
|
||||
success = cv2.imwrite(output_path, image, [cv2.IMWRITE_PNG_COMPRESSION, 3])
|
||||
else:
|
||||
# 如果是单通道，转换为3通道
|
||||
if len(image.shape) == 2:
|
||||
image = cv2.cvtColor(image, cv2.COLOR_GRAY2BGR)
|
||||
success = cv2.imwrite(output_path, image, [cv2.IMWRITE_PNG_COMPRESSION, 3])
|
||||
else:
|
||||
# JPEG格式，确保是3通道BGR
|
||||
if image.shape[2] == 4:
|
||||
# 有alpha通道，先合成到白色背景
|
||||
bgr = image[:, :, :3]
|
||||
alpha = image[:, :, 3:4] / 255.0
|
||||
white_bg = np.ones_like(bgr) * 255
|
||||
image = (bgr * alpha + white_bg * (1 - alpha)).astype(np.uint8)
|
||||
elif len(image.shape) == 2:
|
||||
image = cv2.cvtColor(image, cv2.COLOR_GRAY2BGR)
|
||||
success = cv2.imwrite(output_path, image, [cv2.IMWRITE_JPEG_QUALITY, quality])
|
||||
|
||||
if not success:
|
||||
raise ValueError(f"保存图像失败: {output_path}")
|
||||
elif PIL_AVAILABLE and hasattr(image, 'save'):
|
||||
# 使用PIL保存
|
||||
if is_png:
|
||||
image.save(output_path, 'PNG', compress_level=3)
|
||||
else:
|
||||
# JPEG格式，确保是RGB
|
||||
if image.mode == 'RGBA':
|
||||
# 合成到白色背景
|
||||
rgb = Image.new('RGB', image.size, (255, 255, 255))
|
||||
rgb.paste(image, mask=image.split()[3])
|
||||
image = rgb
|
||||
image.save(output_path, 'JPEG', quality=quality)
|
||||
else:
|
||||
raise ValueError("没有可用的图像保存方法")
|
||||
|
||||
logger.info(f"图像已保存: {output_path}")
|
||||
except Exception as e:
|
||||
logger.error(f"保存图像失败: {e}")
|
||||
raise
|
||||
|
||||
|
||||
def format_time(seconds: float) -> str:
|
||||
"""格式化时间为 HH:MM:SS.mmm"""
|
||||
hours = int(seconds // 3600)
|
||||
minutes = int((seconds % 3600) // 60)
|
||||
secs = int(seconds % 60)
|
||||
milliseconds = int((seconds % 1) * 1000)
|
||||
|
||||
return "02d"
|
||||
|
||||
|
||||
def create_temp_dir(prefix: str = "subtitle_cover_") -> str:
|
||||
"""创建临时目录"""
|
||||
import tempfile
|
||||
temp_dir = tempfile.mkdtemp(prefix=prefix)
|
||||
logger.info(f"创建临时目录: {temp_dir}")
|
||||
return temp_dir
|
||||
|
||||
|
||||
def cleanup_temp_dir(temp_dir: str) -> None:
|
||||
"""清理临时目录"""
|
||||
try:
|
||||
import shutil
|
||||
if os.path.exists(temp_dir):
|
||||
shutil.rmtree(temp_dir)
|
||||
logger.info(f"清理临时目录: {temp_dir}")
|
||||
except Exception as e:
|
||||
logger.warning(f"清理临时目录失败: {e}")
|
||||
|
||||
|
||||
def validate_file_exists(file_path: str, file_type: str = "文件") -> None:
|
||||
"""验证文件是否存在"""
|
||||
if not os.path.exists(file_path):
|
||||
raise FileNotFoundError(f"{file_type}不存在: {file_path}")
|
||||
|
||||
if not os.path.isfile(file_path):
|
||||
raise ValueError(f"{file_type}不是文件: {file_path}")
|
||||
|
||||
|
||||
def get_file_size_mb(file_path: str) -> float:
|
||||
"""获取文件大小(MB)"""
|
||||
size_bytes = os.path.getsize(file_path)
|
||||
return size_bytes / (1024 * 1024)
|
||||
|
||||
|
||||
def calculate_aspect_ratio(width: int, height: int) -> float:
|
||||
"""计算宽高比"""
|
||||
return width / height if height > 0 else 0
|
||||
|
||||
|
||||
def resize_image(image: np.ndarray, target_width: int, target_height: int,
|
||||
keep_aspect_ratio: bool = True) -> np.ndarray:
|
||||
"""调整图像大小"""
|
||||
if keep_aspect_ratio:
|
||||
# 保持宽高比
|
||||
h, w = image.shape[:2]
|
||||
aspect_ratio = w / h
|
||||
|
||||
if target_width / target_height > aspect_ratio:
|
||||
# 目标更宽，以高度为准
|
||||
new_width = int(target_height * aspect_ratio)
|
||||
new_height = target_height
|
||||
else:
|
||||
# 目标更高，以宽度为准
|
||||
new_width = target_width
|
||||
new_height = int(target_width / aspect_ratio)
|
||||
|
||||
resized = cv2.resize(image, (new_width, new_height))
|
||||
else:
|
||||
# 不保持宽高比，直接缩放
|
||||
resized = cv2.resize(image, (target_width, target_height))
|
||||
|
||||
return resized
|
||||
|
||||
|
||||
def blend_images(background: np.ndarray, foreground: np.ndarray,
|
||||
position: Tuple[int, int] = (0, 0)) -> np.ndarray:
|
||||
"""将前景图像合成到背景图像上"""
|
||||
x, y = position
|
||||
h, w = foreground.shape[:2]
|
||||
|
||||
# 确保位置不超出边界
|
||||
bg_h, bg_w = background.shape[:2]
|
||||
x = max(0, min(x, bg_w - w))
|
||||
y = max(0, min(y, bg_h - h))
|
||||
|
||||
# 创建ROI
|
||||
roi = background[y:y+h, x:x+w]
|
||||
|
||||
# 如果前景有alpha通道，进行透明合成
|
||||
if foreground.shape[2] == 4:
|
||||
# 分离颜色和alpha通道
|
||||
foreground_rgb = foreground[:, :, :3]
|
||||
alpha = foreground[:, :, 3] / 255.0
|
||||
|
||||
# 扩展alpha到3通道
|
||||
alpha = np.stack([alpha] * 3, axis=2)
|
||||
|
||||
# 透明合成
|
||||
blended = foreground_rgb * alpha + roi * (1 - alpha)
|
||||
background[y:y+h, x:x+w] = blended.astype(np.uint8)
|
||||
else:
|
||||
# 直接覆盖
|
||||
background[y:y+h, x:x+w] = foreground
|
||||
|
||||
return background
|
||||
|
||||
|
||||
def apply_blur(image: np.ndarray, kernel_size: int = 15) -> np.ndarray:
|
||||
"""应用模糊效果"""
|
||||
return cv2.GaussianBlur(image, (kernel_size, kernel_size), 0)
|
||||
|
||||
|
||||
def apply_gradient_overlay(image: np.ndarray, color: Tuple[int, int, int],
|
||||
opacity: float = 0.5) -> np.ndarray:
|
||||
"""应用渐变覆盖"""
|
||||
overlay = np.full_like(image, color, dtype=np.uint8)
|
||||
return cv2.addWeighted(image, 1 - opacity, overlay, opacity, 0)
|
||||
@@ -0,0 +1,66 @@
|
||||
{
|
||||
"templates": [
|
||||
{
|
||||
"name": "简约风格",
|
||||
"description": "简洁大方的黑色半透明背景，适合各类视频",
|
||||
"config": {
|
||||
"style": "default",
|
||||
"fontSize": 80,
|
||||
"fontColor": [255, 255, 255],
|
||||
"fontName": "NotoSerifCJK-VF", "fontPath": "",
|
||||
"position": "top",
|
||||
"maxWidth": 1920
|
||||
}
|
||||
},
|
||||
{
|
||||
"name": "描边风格",
|
||||
"description": "白色文字配黑色描边，清晰醒目",
|
||||
"config": {
|
||||
"style": "outline",
|
||||
"fontSize": 90,
|
||||
"fontColor": [255, 255, 255],
|
||||
"outlineColor": [0, 0, 0],
|
||||
"fontName": "NotoSerifCJK-VF", "fontPath": "",
|
||||
"position": "top",
|
||||
"maxWidth": 1920
|
||||
}
|
||||
},
|
||||
{
|
||||
"name": "模糊背景",
|
||||
"description": "背景模糊效果，突出标题文字",
|
||||
"config": {
|
||||
"style": "blur_bg",
|
||||
"fontSize": 85,
|
||||
"fontColor": [255, 255, 255],
|
||||
"fontName": "NotoSerifCJK-VF", "fontPath": "",
|
||||
"position": "center",
|
||||
"maxWidth": 1920
|
||||
}
|
||||
},
|
||||
{
|
||||
"name": "渐变遮罩",
|
||||
"description": "优雅的渐变半透明遮罩效果",
|
||||
"config": {
|
||||
"style": "gradient",
|
||||
"fontSize": 75,
|
||||
"fontColor": [255, 255, 255],
|
||||
"fontName": "NotoSerifCJK-VF", "fontPath": "",
|
||||
"position": "top",
|
||||
"maxWidth": 1920
|
||||
}
|
||||
},
|
||||
{
|
||||
"name": "分栏设计",
|
||||
"description": "左侧标题配右侧装饰，设计感十足",
|
||||
"config": {
|
||||
"style": "split",
|
||||
"fontSize": 70,
|
||||
"fontColor": [255, 255, 255],
|
||||
"fontName": "NotoSerifCJK-VF", "fontPath": "",
|
||||
"position": "top",
|
||||
"accentText": "✨",
|
||||
"maxWidth": 1920
|
||||
}
|
||||
}
|
||||
]
|
||||
}
|
||||
@@ -0,0 +1,18 @@
|
||||
# Python 依赖包列表
|
||||
# 用于人像抠图和封面生成功能
|
||||
|
||||
# 核心依赖
|
||||
opencv-python>=4.8.0 # 图像处理
|
||||
numpy>=1.24.0 # 数值计算
|
||||
Pillow>=10.0.0 # 图像处理库
|
||||
loguru>=0.7.0 # 日志记录库（必需）
|
||||
|
||||
# 人像抠图（必需）
|
||||
rembg>=2.0.68 # 背景移除库（包含 U2-Net 模型）
|
||||
onnxruntime>=1.16.0 # ONNX 模型运行时（rembg 依赖）
|
||||
|
||||
# 中文 NLP（必需）
|
||||
jieba>=0.42.1 # 中文分词和关键词提取
|
||||
|
||||
# 可选依赖（提升性能）
|
||||
# onnxruntime-gpu>=1.16.0 # GPU 加速版本（如果有 NVIDIA GPU）
|
||||
@@ -0,0 +1,530 @@
|
||||
#!/usr/bin/env python3
|
||||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
字幕封面生成器 - 一键生成视频字幕和封面
|
||||
Subtitle & Cover Generator for AIGC Panel
|
||||
"""
|
||||
|
||||
import os
|
||||
import sys
|
||||
import json
|
||||
import argparse
|
||||
from typing import Dict, List, Any, Tuple
|
||||
from pathlib import Path
|
||||
|
||||
# 导入必要的库
|
||||
try:
|
||||
import cv2
|
||||
import numpy as np
|
||||
from PIL import Image, ImageDraw, ImageFont, ImageFilter, ImageEnhance, ImageOps
|
||||
from rembg import remove
|
||||
except ImportError as e:
|
||||
print(f"Error: Missing required library - {e}", file=sys.stderr)
|
||||
sys.exit(1)
|
||||
|
||||
|
||||
class SubtitleCoverGenerator:
|
||||
"""字幕封面生成器主类"""
|
||||
|
||||
def __init__(self, config: Dict[str, Any]):
|
||||
"""
|
||||
初始化生成器
|
||||
|
||||
Args:
|
||||
config: 配置字典，包含视频路径、文案、样式等
|
||||
"""
|
||||
self.config = config
|
||||
self.video_path = config.get('videoPath')
|
||||
self.script_text = config.get('scriptText', '')
|
||||
self.subtitle_style = config.get('subtitleStyle', 'default')
|
||||
self.cover_style = config.get('coverStyle', 'default')
|
||||
self.output_dir = config.get('outputDir', './output')
|
||||
|
||||
# 新增配置参数
|
||||
self.is_portrait = config.get('isPortrait', False) # 是否为竖屏视频
|
||||
self.cover_text = config.get('coverText', '') # 封面文字
|
||||
self.cover_text_position = config.get('coverTextPosition', 'bottom') # 文字位置: top, center, bottom
|
||||
self.cover_text_color = config.get('coverTextColor', '#FFFFFF') # 文字颜色
|
||||
self.cover_text_size = config.get('coverTextSize', 50) # 文字大小
|
||||
self.outline_color = config.get('outlineColor', '#FFD700') # 描边颜色
|
||||
self.outline_width = config.get('outlineWidth', 3) # 描边宽度
|
||||
self.blur_radius = config.get('blurRadius', 15) # 背景模糊度
|
||||
self.zoom_level = config.get('zoomLevel', 1.0) # 封面放大倍数
|
||||
|
||||
# 确保输出目录存在
|
||||
os.makedirs(self.output_dir, exist_ok=True)
|
||||
|
||||
def get_video_info(self) -> Dict[str, Any]:
|
||||
"""
|
||||
获取视频信息
|
||||
|
||||
Returns:
|
||||
包含视频宽度、高度、时长等信息的字典
|
||||
"""
|
||||
try:
|
||||
cap = cv2.VideoCapture(self.video_path)
|
||||
if not cap.isOpened():
|
||||
raise Exception("Could not open video file")
|
||||
|
||||
width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
|
||||
height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
|
||||
fps = cap.get(cv2.CAP_PROP_FPS)
|
||||
frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
|
||||
duration = frame_count / fps if fps > 0 else 0
|
||||
|
||||
cap.release()
|
||||
|
||||
# 检测是否为竖屏视频
|
||||
is_portrait = height > width
|
||||
|
||||
return {
|
||||
'width': width,
|
||||
'height': height,
|
||||
'fps': fps,
|
||||
'duration': duration,
|
||||
'is_portrait': is_portrait
|
||||
}
|
||||
except Exception as e:
|
||||
raise Exception(f"Error getting video info: {str(e)}")
|
||||
|
||||
def extract_frame(self, timestamp: float = 1.0) -> np.ndarray:
|
||||
"""
|
||||
从视频中提取关键帧
|
||||
|
||||
Args:
|
||||
timestamp: 提取帧的时间戳（秒）
|
||||
|
||||
Returns:
|
||||
提取的帧（numpy数组）
|
||||
"""
|
||||
try:
|
||||
cap = cv2.VideoCapture(self.video_path)
|
||||
|
||||
# 设置到指定时间戳
|
||||
fps = cap.get(cv2.CAP_PROP_FPS)
|
||||
frame_number = int(timestamp * fps)
|
||||
cap.set(cv2.CAP_PROP_POS_FRAMES, frame_number)
|
||||
|
||||
ret, frame = cap.read()
|
||||
cap.release()
|
||||
|
||||
if not ret:
|
||||
raise Exception("Failed to extract frame from video")
|
||||
|
||||
return frame
|
||||
except Exception as e:
|
||||
raise Exception(f"Error extracting frame: {str(e)}")
|
||||
|
||||
def remove_background(self, image: np.ndarray) -> Image.Image:
|
||||
"""
|
||||
使用rembg进行人物抠图
|
||||
|
||||
Args:
|
||||
image: 输入图像（numpy数组）
|
||||
|
||||
Returns:
|
||||
抠图后的PIL Image对象（带透明通道）
|
||||
"""
|
||||
try:
|
||||
# 转换为PIL Image
|
||||
pil_image = Image.fromarray(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
|
||||
|
||||
# 使用rembg移除背景
|
||||
output = remove(pil_image)
|
||||
|
||||
return output
|
||||
except Exception as e:
|
||||
raise Exception(f"Error removing background: {str(e)}")
|
||||
|
||||
def generate_cover(self) -> str:
|
||||
"""
|
||||
生成视频封面
|
||||
|
||||
Returns:
|
||||
生成的封面图片路径
|
||||
"""
|
||||
try:
|
||||
# 获取视频信息
|
||||
video_info = self.get_video_info()
|
||||
self.is_portrait = video_info['is_portrait']
|
||||
|
||||
# 提取关键帧
|
||||
frame = self.extract_frame(1.0)
|
||||
|
||||
if self.cover_style == 'default':
|
||||
# 默认封面：直接使用提取的帧
|
||||
cover_path = os.path.join(self.output_dir, 'cover.png')
|
||||
cv2.imwrite(cover_path, frame)
|
||||
|
||||
elif self.cover_style == 'blur-bg':
|
||||
# 背景虚化效果
|
||||
cover_path = self._generate_blur_cover(frame)
|
||||
|
||||
elif self.cover_style == 'outline':
|
||||
# 人物描边效果
|
||||
cover_path = self._generate_outline_cover(frame)
|
||||
|
||||
elif self.cover_style == 'gradient':
|
||||
# 渐变背景效果
|
||||
cover_path = self._generate_gradient_cover(frame)
|
||||
|
||||
elif self.cover_style == 'split':
|
||||
# 斜切设计效果
|
||||
cover_path = self._generate_split_cover(frame)
|
||||
else:
|
||||
cover_path = os.path.join(self.output_dir, 'cover.png')
|
||||
cv2.imwrite(cover_path, frame)
|
||||
|
||||
# 如果需要放大封面
|
||||
if self.zoom_level != 1.0:
|
||||
cover_path = self._zoom_cover(cover_path)
|
||||
|
||||
# 如果有文字需要添加
|
||||
if self.cover_text.strip():
|
||||
cover_path = self._add_text_to_cover(cover_path)
|
||||
|
||||
return cover_path
|
||||
except Exception as e:
|
||||
raise Exception(f"Error generating cover: {str(e)}")
|
||||
|
||||
def _zoom_cover(self, cover_path: str) -> str:
|
||||
"""放大封面"""
|
||||
try:
|
||||
image = Image.open(cover_path)
|
||||
|
||||
# 计算新的尺寸
|
||||
new_width = int(image.width * self.zoom_level)
|
||||
new_height = int(image.height * self.zoom_level)
|
||||
|
||||
# 放大图像
|
||||
zoomed = image.resize((new_width, new_height), Image.Resampling.LANCZOS)
|
||||
|
||||
# 如果是竖屏视频，需要裁剪中心区域
|
||||
if self.is_portrait:
|
||||
# 裁剪为原始尺寸的正方形区域
|
||||
left = (zoomed.width - image.width) // 2
|
||||
top = (zoomed.height - image.height) // 2
|
||||
right = left + image.width
|
||||
bottom = top + image.height
|
||||
zoomed = zoomed.crop((left, top, right, bottom))
|
||||
else:
|
||||
# 横屏视频裁剪中心区域
|
||||
left = (zoomed.width - image.width) // 2
|
||||
top = (zoomed.height - image.height) // 2
|
||||
right = left + image.width
|
||||
bottom = top + image.height
|
||||
zoomed = zoomed.crop((left, top, right, bottom))
|
||||
|
||||
# 保存放大后的封面
|
||||
zoom_path = os.path.join(self.output_dir, 'cover_zoomed.png')
|
||||
zoomed.save(zoom_path, quality=95)
|
||||
|
||||
return zoom_path
|
||||
except Exception as e:
|
||||
print(f"Warning: Could not zoom cover: {e}", file=sys.stderr)
|
||||
return cover_path
|
||||
|
||||
def _add_text_to_cover(self, cover_path: str) -> str:
|
||||
"""添加文字到封面"""
|
||||
try:
|
||||
image = Image.open(cover_path)
|
||||
draw = ImageDraw.Draw(image)
|
||||
|
||||
# 尝试加载字体，如果失败则使用默认字体
|
||||
try:
|
||||
font = ImageFont.truetype("arial.ttf", self.cover_text_size)
|
||||
except:
|
||||
try:
|
||||
font = ImageFont.truetype("DejaVuSans-Bold.ttf", self.cover_text_size)
|
||||
except:
|
||||
font = ImageFont.load_default()
|
||||
|
||||
# 计算文字位置
|
||||
bbox = draw.textbbox((0, 0), self.cover_text, font=font)
|
||||
text_width = bbox[2] - bbox[0]
|
||||
text_height = bbox[3] - bbox[1]
|
||||
|
||||
if self.cover_text_position == 'top':
|
||||
x = (image.width - text_width) // 2
|
||||
y = 50
|
||||
elif self.cover_text_position == 'center':
|
||||
x = (image.width - text_width) // 2
|
||||
y = (image.height - text_height) // 2
|
||||
else: # bottom
|
||||
x = (image.width - text_width) // 2
|
||||
y = image.height - text_height - 50
|
||||
|
||||
# 添加描边效果
|
||||
if self.outline_width > 0:
|
||||
# 创建描边
|
||||
outline_color = self._hex_to_rgb(self.outline_color)
|
||||
for offset_x in range(-self.outline_width, self.outline_width + 1):
|
||||
for offset_y in range(-self.outline_width, self.outline_width + 1):
|
||||
if offset_x == 0 and offset_y == 0:
|
||||
continue
|
||||
draw.text((x + offset_x, y + offset_y), self.cover_text, fill=outline_color, font=font)
|
||||
|
||||
# 添加文字
|
||||
text_color = self._hex_to_rgb(self.cover_text_color)
|
||||
draw.text((x, y), self.cover_text, fill=text_color, font=font)
|
||||
|
||||
# 保存带文字的封面
|
||||
text_path = os.path.join(self.output_dir, 'cover_with_text.png')
|
||||
image.save(text_path, quality=95)
|
||||
|
||||
return text_path
|
||||
except Exception as e:
|
||||
print(f"Warning: Could not add text to cover: {e}", file=sys.stderr)
|
||||
return cover_path
|
||||
|
||||
def _hex_to_rgb(self, hex_color: str) -> Tuple[int, int, int]:
|
||||
"""将十六进制颜色转换为RGB元组"""
|
||||
hex_color = hex_color.lstrip('#')
|
||||
return tuple(int(hex_color[i:i+2], 16) for i in (0, 2, 4))
|
||||
|
||||
def _generate_blur_cover(self, frame: np.ndarray) -> str:
|
||||
"""生成背景虚化封面"""
|
||||
# 转换为PIL
|
||||
image = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
|
||||
|
||||
# 抠出人物
|
||||
person = self.remove_background(frame)
|
||||
|
||||
# 背景虚化（使用配置的模糊度）
|
||||
background = image.filter(ImageFilter.GaussianBlur(radius=self.blur_radius))
|
||||
|
||||
# 合成
|
||||
background.paste(person, (0, 0), person)
|
||||
|
||||
# 保存
|
||||
cover_path = os.path.join(self.output_dir, 'cover_blur.png')
|
||||
background.save(cover_path, quality=95)
|
||||
|
||||
return cover_path
|
||||
|
||||
def _generate_outline_cover(self, frame: np.ndarray) -> str:
|
||||
"""生成人物描边封面"""
|
||||
# 转换为PIL
|
||||
image = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
|
||||
|
||||
# 抠出人物
|
||||
person = self.remove_background(frame)
|
||||
|
||||
# 创建描边效果
|
||||
# 将person转为numpy进行描边处理
|
||||
person_np = np.array(person)
|
||||
alpha = person_np[:, :, 3]
|
||||
|
||||
# 使用形态学操作创建描边
|
||||
kernel = np.ones((self.outline_width * 2 + 1, self.outline_width * 2 + 1), np.uint8)
|
||||
dilated = cv2.dilate(alpha, kernel, iterations=1)
|
||||
outline = dilated - alpha
|
||||
|
||||
# 创建描边图层
|
||||
outline_layer = Image.new('RGBA', person.size, (0, 0, 0, 0))
|
||||
outline_draw = ImageDraw.Draw(outline_layer)
|
||||
|
||||
# 获取描边颜色
|
||||
outline_color = self._hex_to_rgb(self.outline_color)
|
||||
|
||||
# 应用描边
|
||||
for y in range(outline.shape[0]):
|
||||
for x in range(outline.shape[1]):
|
||||
if outline[y, x] > 0:
|
||||
outline_layer.putpixel((x, y), outline_color + (255,))
|
||||
|
||||
# 合成
|
||||
result = Image.new('RGBA', image.size)
|
||||
result.paste(image, (0, 0))
|
||||
result.paste(outline_layer, (0, 0), outline_layer)
|
||||
result.paste(person, (0, 0), person)
|
||||
|
||||
# 保存
|
||||
cover_path = os.path.join(self.output_dir, 'cover_outline.png')
|
||||
result.save(cover_path, quality=95)
|
||||
|
||||
return cover_path
|
||||
|
||||
def _generate_gradient_cover(self, frame: np.ndarray) -> str:
|
||||
"""生成渐变背景封面"""
|
||||
# 转换为PIL
|
||||
image = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
|
||||
|
||||
# 抠出人物
|
||||
person = self.remove_background(frame)
|
||||
|
||||
# 创建渐变背景
|
||||
width, height = image.size
|
||||
gradient = Image.new('RGB', (width, height), color=(0, 0, 0))
|
||||
draw = ImageDraw.Draw(gradient)
|
||||
|
||||
# 绘制渐变
|
||||
for y in range(height):
|
||||
r = int(103 + (118 - 103) * y / height)
|
||||
g = int(126 + (74 - 126) * y / height)
|
||||
b = int(234 + (162 - 234) * y / height)
|
||||
draw.line([(0, y), (width, y)], fill=(r, g, b))
|
||||
|
||||
# 合成
|
||||
gradient.paste(person, (0, 0), person)
|
||||
|
||||
# 保存
|
||||
cover_path = os.path.join(self.output_dir, 'cover_gradient.png')
|
||||
gradient.save(cover_path, quality=95)
|
||||
|
||||
return cover_path
|
||||
|
||||
def _generate_split_cover(self, frame: np.ndarray) -> str:
|
||||
"""生成斜切设计封面"""
|
||||
# 转换为PIL
|
||||
image = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
|
||||
width, height = image.size
|
||||
|
||||
# 创建斜切背景
|
||||
background = Image.new('RGB', (width, height), color=(255, 255, 255))
|
||||
draw = ImageDraw.Draw(background)
|
||||
|
||||
# 绘制斜切色块
|
||||
draw.polygon([(0, 0), (width * 0.6, 0), (width * 0.4, height), (0, height)],
|
||||
fill=(103, 126, 234))
|
||||
draw.polygon([(width * 0.6, 0), (width, 0), (width, height), (width * 0.4, height)],
|
||||
fill=(118, 74, 162))
|
||||
|
||||
# 应用半透明图片
|
||||
image_enhanced = ImageEnhance.Brightness(image).enhance(0.8)
|
||||
background.paste(image_enhanced, (0, 0))
|
||||
|
||||
# 保存
|
||||
cover_path = os.path.join(self.output_dir, 'cover_split.png')
|
||||
background.save(cover_path, quality=95)
|
||||
|
||||
return cover_path
|
||||
|
||||
def generate_subtitles(self) -> List[Dict[str, Any]]:
|
||||
"""
|
||||
生成字幕列表
|
||||
|
||||
Returns:
|
||||
字幕列表，每个字幕包含文本、开始时间、结束时间
|
||||
"""
|
||||
try:
|
||||
# 将文案按行分割
|
||||
lines = [line.strip() for line in self.script_text.split('\n') if line.strip()]
|
||||
|
||||
# 生成字幕列表（简单示例，每句3秒）
|
||||
subtitles = []
|
||||
for i, line in enumerate(lines):
|
||||
subtitle = {
|
||||
'index': i + 1,
|
||||
'text': line,
|
||||
'start_time': i * 3.0,
|
||||
'end_time': (i + 1) * 3.0,
|
||||
'style': self.subtitle_style
|
||||
}
|
||||
subtitles.append(subtitle)
|
||||
|
||||
return subtitles
|
||||
except Exception as e:
|
||||
raise Exception(f"Error generating subtitles: {str(e)}")
|
||||
|
||||
def save_srt(self, subtitles: List[Dict[str, Any]]) -> str:
|
||||
"""
|
||||
保存SRT字幕文件
|
||||
|
||||
Args:
|
||||
subtitles: 字幕列表
|
||||
|
||||
Returns:
|
||||
SRT文件路径
|
||||
"""
|
||||
try:
|
||||
srt_path = os.path.join(self.output_dir, 'subtitles.srt')
|
||||
|
||||
with open(srt_path, 'w', encoding='utf-8') as f:
|
||||
for sub in subtitles:
|
||||
# SRT格式
|
||||
f.write(f"{sub['index']}\n")
|
||||
f.write(f"{self._format_time(sub['start_time'])} --> {self._format_time(sub['end_time'])}\n")
|
||||
f.write(f"{sub['text']}\n\n")
|
||||
|
||||
return srt_path
|
||||
except Exception as e:
|
||||
raise Exception(f"Error saving SRT file: {str(e)}")
|
||||
|
||||
def _format_time(self, seconds: float) -> str:
|
||||
"""格式化时间为SRT格式"""
|
||||
hours = int(seconds // 3600)
|
||||
minutes = int((seconds % 3600) // 60)
|
||||
secs = int(seconds % 60)
|
||||
millis = int((seconds % 1) * 1000)
|
||||
|
||||
return f"{hours:02d}:{minutes:02d}:{secs:02d},{millis:03d}"
|
||||
|
||||
def generate_all(self) -> Dict[str, Any]:
|
||||
"""
|
||||
一键生成所有内容
|
||||
|
||||
Returns:
|
||||
包含生成结果的字典
|
||||
"""
|
||||
try:
|
||||
# 生成封面
|
||||
print("Generating cover...", file=sys.stderr)
|
||||
cover_path = self.generate_cover()
|
||||
|
||||
# 生成字幕
|
||||
print("Generating subtitles...", file=sys.stderr)
|
||||
subtitles = self.generate_subtitles()
|
||||
srt_path = self.save_srt(subtitles)
|
||||
|
||||
result = {
|
||||
'success': True,
|
||||
'coverPath': cover_path,
|
||||
'srtPath': srt_path,
|
||||
'subtitles': subtitles,
|
||||
'message': 'Generated successfully'
|
||||
}
|
||||
|
||||
return result
|
||||
except Exception as e:
|
||||
return {
|
||||
'success': False,
|
||||
'error': str(e),
|
||||
'message': f'Generation failed: {str(e)}'
|
||||
}
|
||||
|
||||
|
||||
def main():
|
||||
"""主函数"""
|
||||
parser = argparse.ArgumentParser(description='Subtitle & Cover Generator')
|
||||
parser.add_argument('config_file', help='Configuration JSON file path')
|
||||
args = parser.parse_args()
|
||||
|
||||
try:
|
||||
# 读取配置文件
|
||||
with open(args.config_file, 'r', encoding='utf-8') as f:
|
||||
config = json.load(f)
|
||||
|
||||
# 创建生成器
|
||||
generator = SubtitleCoverGenerator(config)
|
||||
|
||||
# 执行生成
|
||||
result = generator.generate_all()
|
||||
|
||||
# 输出结果
|
||||
print(json.dumps(result, ensure_ascii=False))
|
||||
|
||||
sys.exit(0 if result['success'] else 1)
|
||||
|
||||
except Exception as e:
|
||||
error_result = {
|
||||
'success': False,
|
||||
'error': str(e),
|
||||
'message': f'Error: {str(e)}'
|
||||
}
|
||||
print(json.dumps(error_result, ensure_ascii=False))
|
||||
sys.exit(1)
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
main()
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,88 @@
|
||||
#!/usr/bin/env python
|
||||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
Simple Python dependency test script
|
||||
Can be run with full path, e.g.:
|
||||
C:\Python312\python.exe python/test_dependencies_simple.py
|
||||
"""
|
||||
|
||||
import sys
|
||||
import os
|
||||
|
||||
def test_python_version():
|
||||
"""Test Python version"""
|
||||
version = sys.version_info
|
||||
print(f"Python version: {version.major}.{version.minor}.{version.micro}")
|
||||
|
||||
if version.major < 3 or (version.major == 3 and version.minor < 7):
|
||||
print("X Python version too low, need 3.7 or higher")
|
||||
return False
|
||||
else:
|
||||
print("[OK] Python version meets requirements")
|
||||
return True
|
||||
|
||||
def test_pillow():
|
||||
"""Test Pillow library"""
|
||||
try:
|
||||
from PIL import Image
|
||||
import PIL
|
||||
print(f"[OK] Pillow installed: {PIL.__version__}")
|
||||
return True
|
||||
except ImportError:
|
||||
print("X Pillow not installed")
|
||||
print(" Install command: pip install Pillow")
|
||||
return False
|
||||
|
||||
def test_numpy():
|
||||
"""Test NumPy library"""
|
||||
try:
|
||||
import numpy as np
|
||||
print(f"[OK] NumPy installed: {np.__version__}")
|
||||
return True
|
||||
except ImportError:
|
||||
print("X NumPy not installed")
|
||||
print(" Install command: pip install numpy")
|
||||
return False
|
||||
|
||||
def main():
|
||||
print("=" * 50)
|
||||
print("Python Dependency Test")
|
||||
print("=" * 50)
|
||||
print()
|
||||
|
||||
print(f"Python executable: {sys.executable}")
|
||||
print()
|
||||
|
||||
results = []
|
||||
|
||||
# Test Python version
|
||||
results.append(test_python_version())
|
||||
print()
|
||||
|
||||
# Test Pillow
|
||||
results.append(test_pillow())
|
||||
print()
|
||||
|
||||
# Test NumPy
|
||||
results.append(test_numpy())
|
||||
print()
|
||||
|
||||
# Summary
|
||||
print("=" * 50)
|
||||
if all(results):
|
||||
print("[OK] All dependencies ready!")
|
||||
print()
|
||||
print("Next steps:")
|
||||
print("1. Add Python to system PATH, or")
|
||||
print("2. Use full path in code:")
|
||||
print(f" {sys.executable}")
|
||||
return 0
|
||||
else:
|
||||
print("X Missing dependencies")
|
||||
print()
|
||||
print("Install all dependencies:")
|
||||
print(" pip install Pillow numpy")
|
||||
return 1
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(main())
|
||||
@@ -0,0 +1,68 @@
|
||||
#!/usr/bin/env python3
|
||||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
测试Python环境和依赖库
|
||||
Test Python Environment and Dependencies
|
||||
"""
|
||||
|
||||
import sys
|
||||
|
||||
|
||||
def check_library(lib_name, import_name=None):
|
||||
"""检查库是否可用"""
|
||||
if import_name is None:
|
||||
import_name = lib_name
|
||||
|
||||
try:
|
||||
__import__(import_name)
|
||||
print(f"[OK] {lib_name} is installed")
|
||||
return True
|
||||
except ImportError:
|
||||
print(f"[ERROR] {lib_name} is NOT installed")
|
||||
return False
|
||||
|
||||
|
||||
def main():
|
||||
"""主函数"""
|
||||
print("=" * 50)
|
||||
print("Python Environment Check")
|
||||
print("=" * 50)
|
||||
print(f"Python version: {sys.version}")
|
||||
print()
|
||||
|
||||
# 检查必要的库
|
||||
libraries = [
|
||||
("OpenCV", "cv2"),
|
||||
("NumPy", "numpy"),
|
||||
("Pillow", "PIL"),
|
||||
("rembg", "rembg"),
|
||||
("ffmpeg-python", "ffmpeg"),
|
||||
("spacy", "spacy"),
|
||||
("transformers", "transformers"),
|
||||
]
|
||||
|
||||
print("Checking required libraries:")
|
||||
print("-" * 50)
|
||||
|
||||
all_installed = True
|
||||
for lib_name, import_name in libraries:
|
||||
if not check_library(lib_name, import_name):
|
||||
all_installed = False
|
||||
|
||||
print()
|
||||
print("=" * 50)
|
||||
if all_installed:
|
||||
print("[SUCCESS] All dependencies are installed!")
|
||||
print("Environment is ready for subtitle and cover generation.")
|
||||
else:
|
||||
print("[FAILED] Some dependencies are missing.")
|
||||
print("Please install missing libraries using:")
|
||||
print("pip install opencv-python numpy Pillow rembg ffmpeg-python spacy transformers")
|
||||
|
||||
print("=" * 50)
|
||||
|
||||
return 0 if all_installed else 1
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
sys.exit(main())
|
||||
@@ -0,0 +1,71 @@
|
||||
#!/usr/bin/env python3
|
||||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
测试完整字幕和封面生成
|
||||
"""
|
||||
|
||||
import os
|
||||
import json
|
||||
import sys
|
||||
|
||||
# 配置信息
|
||||
CONFIG = {
|
||||
"videoPath": r"C:\Users\Administrator\Desktop\视频生成\对比结果\121212.mp4",
|
||||
"scriptText": "今天天气真好！\n明天下雨。\n我们去公园玩吧。\n这个电影太好看了。\n质量不好，我很失望。",
|
||||
"subtitleStyle": "emotion-mixed", # 使用情感混合模式
|
||||
"coverStyle": "default",
|
||||
"outputDir": r"C:\Users\Administrator\AppData\Roaming\aigcpanel\data\temp\test_full_generation",
|
||||
"useFunasr": False,
|
||||
"coverOnly": False
|
||||
}
|
||||
|
||||
def main():
|
||||
print("=" * 60)
|
||||
print("测试完整字幕和封面生成")
|
||||
print("=" * 60)
|
||||
|
||||
# 确保输出目录存在
|
||||
os.makedirs(CONFIG["outputDir"], exist_ok=True)
|
||||
|
||||
# 保存配置文件
|
||||
config_path = os.path.join(CONFIG["outputDir"], "config.json")
|
||||
with open(config_path, 'w', encoding='utf-8') as f:
|
||||
json.dump(CONFIG, f, ensure_ascii=False, indent=2)
|
||||
|
||||
print(f"\n配置文件已保存: {config_path}")
|
||||
print("\n开始生成...\n")
|
||||
|
||||
# 调用字幕生成器
|
||||
from subtitle_cover_generator_simple import SubtitleCoverGenerator
|
||||
|
||||
generator = SubtitleCoverGenerator(CONFIG)
|
||||
result = generator.generate_all()
|
||||
|
||||
print("\n" + "=" * 60)
|
||||
print("生成结果:")
|
||||
print("=" * 60)
|
||||
print(json.dumps(result, ensure_ascii=False, indent=2))
|
||||
|
||||
if result['success']:
|
||||
print("\n✓ 生成成功!")
|
||||
print(f"\n封面路径: {result.get('coverPath')}")
|
||||
print(f"视频路径: {result.get('videoPath')}")
|
||||
print(f"\nNLP分析结果:")
|
||||
print(f" 标题: {result['nlpAnalysis'].get('title')}")
|
||||
print(f" 关键词: {result['nlpAnalysis'].get('keywords')}")
|
||||
print(f" 情感: {result['nlpAnalysis'].get('sentiment')} (分数: {result['nlpAnalysis'].get('sentiment_score', 0):.4f})")
|
||||
print(f"\n字幕数量: {len(result.get('subtitles', []))}")
|
||||
|
||||
# 打印每个字幕的情感和样式
|
||||
print("\n字幕详情:")
|
||||
for i, subtitle in enumerate(result.get('subtitles', []), 1):
|
||||
print(f" {i}. {subtitle['text']}")
|
||||
print(f" 情感: {subtitle['sentiment']} (分数: {subtitle['sentiment_score']:.4f}) -> 样式: {subtitle['style']}")
|
||||
|
||||
return 0
|
||||
else:
|
||||
print(f"\n✗ 生成失败: {result.get('error')}")
|
||||
return 1
|
||||
|
||||
if __name__ == '__main__':
|
||||
sys.exit(main())
|
||||
@@ -0,0 +1,170 @@
|
||||
#!/usr/bin/env python3
|
||||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
测试字幕封面生成器的核心功能
|
||||
Test core functionality of subtitle and cover generator
|
||||
"""
|
||||
|
||||
import os
|
||||
import sys
|
||||
import json
|
||||
import tempfile
|
||||
from pathlib import Path
|
||||
|
||||
def test_video_frame_extraction():
|
||||
"""测试视频帧提取功能"""
|
||||
try:
|
||||
import cv2
|
||||
print("✓ OpenCV loaded successfully")
|
||||
|
||||
# 创建一个假的视频文件路径来测试
|
||||
# (实际上不会提取帧，只是测试cv2.VideoCapture是否能初始化)
|
||||
cap = cv2.VideoCapture()
|
||||
print("✓ VideoCapture initialized successfully")
|
||||
cap.release()
|
||||
|
||||
return True
|
||||
except Exception as e:
|
||||
print(f"✗ Video frame extraction failed: {e}")
|
||||
return False
|
||||
|
||||
def test_image_processing():
|
||||
"""测试图像处理功能"""
|
||||
try:
|
||||
import numpy as np
|
||||
from PIL import Image, ImageDraw
|
||||
print("✓ PIL and NumPy loaded successfully")
|
||||
|
||||
# 创建一个测试图像
|
||||
img = Image.new('RGB', (100, 100), color='red')
|
||||
draw = ImageDraw.Draw(img)
|
||||
draw.text((10, 10), "Test", fill='white')
|
||||
|
||||
# 转换为numpy数组
|
||||
img_np = np.array(img)
|
||||
print(f"✓ Image processing works: shape {img_np.shape}")
|
||||
|
||||
return True
|
||||
except Exception as e:
|
||||
print(f"✗ Image processing failed: {e}")
|
||||
return False
|
||||
|
||||
def test_nlp_processing():
|
||||
"""测试NLP处理功能"""
|
||||
try:
|
||||
import jieba
|
||||
print("✓ Jieba loaded successfully")
|
||||
|
||||
# 测试关键词提取
|
||||
text = "这是一个测试句子，用于验证中文分词功能"
|
||||
words = jieba.cut(text)
|
||||
keywords = jieba.analyse.extract_tags(text, topK=5)
|
||||
print(f"✓ NLP processing works: keywords {keywords}")
|
||||
|
||||
return True
|
||||
except Exception as e:
|
||||
print(f"✗ NLP processing failed: {e}")
|
||||
return False
|
||||
|
||||
def test_background_removal():
|
||||
"""测试背景移除功能"""
|
||||
try:
|
||||
from rembg import remove
|
||||
print("✓ RMBG loaded successfully")
|
||||
|
||||
# 创建一个简单的测试图像
|
||||
from PIL import Image
|
||||
test_img = Image.new('RGB', (50, 50), color='blue')
|
||||
# 测试是否能调用remove函数（不实际处理）
|
||||
print("✓ Background removal library available")
|
||||
|
||||
return True
|
||||
except Exception as e:
|
||||
print(f"✗ Background removal failed: {e}")
|
||||
return False
|
||||
|
||||
def test_transformers():
|
||||
"""测试Transformers功能"""
|
||||
try:
|
||||
from transformers import pipeline
|
||||
import torch
|
||||
print("✓ Transformers loaded successfully")
|
||||
print(f"✓ PyTorch available: CUDA={torch.cuda.is_available()}")
|
||||
|
||||
# 不实际加载模型，只测试导入
|
||||
return True
|
||||
except Exception as e:
|
||||
print(f"✗ Transformers failed: {e}")
|
||||
return False
|
||||
|
||||
def test_funasr():
|
||||
"""测试FUNASR功能"""
|
||||
try:
|
||||
from funasr import AutoModel
|
||||
print("✓ FUNASR loaded successfully")
|
||||
|
||||
# 测试是否能导入AutoModel
|
||||
return True
|
||||
except Exception as e:
|
||||
print(f"✗ FUNASR failed: {e}")
|
||||
return False
|
||||
|
||||
def test_ffmpeg():
|
||||
"""测试FFmpeg功能"""
|
||||
try:
|
||||
import subprocess
|
||||
result = subprocess.run(['ffmpeg', '-version'],
|
||||
capture_output=True, text=True, timeout=10)
|
||||
if result.returncode == 0:
|
||||
version_line = result.stdout.split('\n')[0]
|
||||
print(f"✓ FFmpeg available: {version_line}")
|
||||
return True
|
||||
else:
|
||||
print("✗ FFmpeg not working properly")
|
||||
return False
|
||||
except Exception as e:
|
||||
print(f"✗ FFmpeg test failed: {e}")
|
||||
return False
|
||||
|
||||
def main():
|
||||
"""主测试函数"""
|
||||
print("=" * 60)
|
||||
print("字幕封面生成器功能测试")
|
||||
print("Subtitle & Cover Generator Feature Test")
|
||||
print("=" * 60)
|
||||
|
||||
tests = [
|
||||
("视频帧提取", test_video_frame_extraction),
|
||||
("图像处理", test_image_processing),
|
||||
("NLP处理", test_nlp_processing),
|
||||
("背景移除", test_background_removal),
|
||||
("Transformers", test_transformers),
|
||||
("FUNASR语音识别", test_funasr),
|
||||
("FFmpeg", test_ffmpeg),
|
||||
]
|
||||
|
||||
passed = 0
|
||||
total = len(tests)
|
||||
|
||||
for test_name, test_func in tests:
|
||||
print(f"\n测试: {test_name}")
|
||||
print("-" * 30)
|
||||
if test_func():
|
||||
passed += 1
|
||||
print()
|
||||
|
||||
print("=" * 60)
|
||||
print(f"测试结果: {passed}/{total} 通过")
|
||||
print(f"Test Results: {passed}/{total} passed")
|
||||
|
||||
if passed == total:
|
||||
print("🎉 所有功能正常！可以开始使用字幕封面生成功能了")
|
||||
print("All features working! Ready to use subtitle and cover generation.")
|
||||
return 0
|
||||
else:
|
||||
print("⚠️ 某些功能不可用，请检查依赖库安装")
|
||||
print("Some features not available, please check dependencies.")
|
||||
return 1
|
||||
|
||||
if __name__ == '__main__':
|
||||
sys.exit(main())
|
||||
@@ -0,0 +1,57 @@
|
||||
#!/usr/bin/env python3
|
||||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
测试脚本：触发 rembg 模型下载
|
||||
"""
|
||||
|
||||
import sys
|
||||
import os
|
||||
from pathlib import Path
|
||||
|
||||
# 设置控制台编码为 UTF-8
|
||||
if sys.platform == 'win32':
|
||||
os.system('chcp 65001 >nul')
|
||||
|
||||
print("=" * 60)
|
||||
print("测试 rembg 模型加载")
|
||||
print("=" * 60)
|
||||
print()
|
||||
|
||||
try:
|
||||
print("[1/3] 导入 rembg...")
|
||||
from rembg import new_session
|
||||
print("[OK] rembg 导入成功")
|
||||
print()
|
||||
|
||||
print("[2/3] 初始化 u2net 模型...")
|
||||
print("注意：首次运行会自动下载模型（约 176MB）")
|
||||
print("模型将缓存到：", Path.home() / ".u2net")
|
||||
print()
|
||||
|
||||
session = new_session('u2net')
|
||||
print("[OK] u2net 模型加载成功！")
|
||||
print()
|
||||
|
||||
print("[3/3] 验证模型...")
|
||||
print(f"Session 类型: {type(session)}")
|
||||
print(f"Session 信息: {session}")
|
||||
print()
|
||||
|
||||
print("=" * 60)
|
||||
print("[OK] 所有测试通过！")
|
||||
print("=" * 60)
|
||||
print()
|
||||
print("模型已准备就绪，可以开始使用人像抠图功能。")
|
||||
print()
|
||||
|
||||
except Exception as e:
|
||||
print()
|
||||
print("=" * 60)
|
||||
print("[ERROR] 测试失败")
|
||||
print("=" * 60)
|
||||
print()
|
||||
print(f"错误信息: {e}")
|
||||
print()
|
||||
import traceback
|
||||
traceback.print_exc()
|
||||
sys.exit(1)
|
||||
@@ -0,0 +1,206 @@
|
||||
#!/usr/bin/env python3
|
||||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
测试情感驱动的字幕生成功能
|
||||
Test Sentiment-Driven Subtitle Generation
|
||||
"""
|
||||
|
||||
import os
|
||||
import sys
|
||||
import json
|
||||
import tempfile
|
||||
|
||||
def create_test_config():
|
||||
"""创建测试配置"""
|
||||
config = {
|
||||
'videoPath': 'test_video.mp4', # 需要替换为实际视频路径
|
||||
'scriptText': """今天真是太开心了！
|
||||
这个产品质量很糟糕。
|
||||
明天我们去公园玩。
|
||||
这个电影太精彩了！
|
||||
天气不好，我很失望。
|
||||
普通的一天，没什么特别的。""",
|
||||
'subtitleStyle': 'default', # 这会被情感分析覆盖
|
||||
'coverStyle': 'default',
|
||||
'outputDir': './test_output',
|
||||
'useFunasr': False, # 使用文本模式
|
||||
'coverOnly': False
|
||||
}
|
||||
|
||||
return config
|
||||
|
||||
def test_sentiment_analysis():
|
||||
"""测试情感分析功能"""
|
||||
print("=" * 60)
|
||||
print("测试情感分析功能")
|
||||
print("=" * 60)
|
||||
|
||||
try:
|
||||
from subtitle_cover_generator_simple import SubtitleCoverGenerator
|
||||
|
||||
config = create_test_config()
|
||||
generator = SubtitleCoverGenerator(config)
|
||||
|
||||
# 测试句子
|
||||
test_sentences = [
|
||||
"今天真是太开心了！", # 应该是positive
|
||||
"这个产品质量很糟糕。", # 应该是negative
|
||||
"明天我们去公园玩。", # 应该是neutral
|
||||
"这个电影太精彩了！", # 应该是positive (强)
|
||||
"天气不好，我很失望。", # 应该是negative
|
||||
]
|
||||
|
||||
print("\n逐句情感分析测试:")
|
||||
print("-" * 60)
|
||||
|
||||
for sentence in test_sentences:
|
||||
result = generator.analyze_sentence_sentiment(sentence)
|
||||
sentiment = result['sentiment']
|
||||
score = result['score']
|
||||
style = generator.get_style_by_sentiment(sentiment, score)
|
||||
|
||||
print(f"\n句子: {sentence}")
|
||||
print(f" 情感: {sentiment} (分数: {score:.4f})")
|
||||
print(f" 推荐样式: {style}")
|
||||
|
||||
print("\n" + "=" * 60)
|
||||
print("✓ 情感分析功能测试通过!")
|
||||
return True
|
||||
|
||||
except Exception as e:
|
||||
print(f"\n✗ 情感分析测试失败: {e}")
|
||||
import traceback
|
||||
traceback.print_exc()
|
||||
return False
|
||||
|
||||
def test_subtitle_generation():
|
||||
"""测试字幕生成功能"""
|
||||
print("\n" + "=" * 60)
|
||||
print("测试字幕生成功能")
|
||||
print("=" * 60)
|
||||
|
||||
try:
|
||||
from subtitle_cover_generator_simple import SubtitleCoverGenerator
|
||||
|
||||
config = create_test_config()
|
||||
generator = SubtitleCoverGenerator(config)
|
||||
|
||||
# 生成字幕
|
||||
subtitles = generator.generate_subtitles()
|
||||
|
||||
print(f"\n生成了 {len(subtitles)} 条字幕:")
|
||||
print("-" * 60)
|
||||
|
||||
for sub in subtitles:
|
||||
print(f"\n[{sub['index']}] {sub['start_time']:.1f}s - {sub['end_time']:.1f}s")
|
||||
print(f" 文本: {sub['text']}")
|
||||
print(f" 情感: {sub.get('sentiment', 'N/A')} (分数: {sub.get('sentiment_score', 0):.4f})")
|
||||
print(f" 样式: {sub.get('style', 'N/A')}")
|
||||
|
||||
print("\n" + "=" * 60)
|
||||
print("✓ 字幕生成功能测试通过!")
|
||||
return True
|
||||
|
||||
except Exception as e:
|
||||
print(f"\n✗ 字幕生成测试失败: {e}")
|
||||
import traceback
|
||||
traceback.print_exc()
|
||||
return False
|
||||
|
||||
def test_ass_file_generation():
|
||||
"""测试ASS字幕文件生成"""
|
||||
print("\n" + "=" * 60)
|
||||
print("测试ASS字幕文件生成")
|
||||
print("=" * 60)
|
||||
|
||||
try:
|
||||
from subtitle_cover_generator_simple import SubtitleCoverGenerator
|
||||
|
||||
config = create_test_config()
|
||||
os.makedirs(config['outputDir'], exist_ok=True)
|
||||
|
||||
generator = SubtitleCoverGenerator(config)
|
||||
|
||||
# 生成字幕
|
||||
subtitles = generator.generate_subtitles()
|
||||
|
||||
# 保存ASS文件
|
||||
ass_path = generator.save_ass(subtitles)
|
||||
|
||||
# 读取并显示ASS文件内容
|
||||
with open(ass_path, 'r', encoding='utf-8-sig') as f:
|
||||
content = f.read()
|
||||
|
||||
print(f"\nASS文件已生成: {ass_path}")
|
||||
print("\n文件内容预览 (前30行):")
|
||||
print("-" * 60)
|
||||
lines = content.split('\n')[:30]
|
||||
for i, line in enumerate(lines, 1):
|
||||
print(f"{i:3d}: {line}")
|
||||
|
||||
print("\n" + "=" * 60)
|
||||
print("✓ ASS字幕文件生成测试通过!")
|
||||
return True
|
||||
|
||||
except Exception as e:
|
||||
print(f"\n✗ ASS文件生成测试失败: {e}")
|
||||
import traceback
|
||||
traceback.print_exc()
|
||||
return False
|
||||
|
||||
def print_summary():
|
||||
"""打印功能说明"""
|
||||
print("\n" + "=" * 60)
|
||||
print("情感驱动字幕功能说明")
|
||||
print("=" * 60)
|
||||
print("""
|
||||
工作原理:
|
||||
1. 对每句话进行情感分析（使用Transformer模型或关键词匹配）
|
||||
2. 根据情感类型和强度自动选择字幕样式:
|
||||
|
||||
情感类型 情感强度 字幕样式 视觉效果
|
||||
--------------------------------------------------------
|
||||
positive > 0.8 neon 霓虹样式-明亮彩色描边
|
||||
positive <= 0.8 highlight 高亮样式-黄色背景
|
||||
negative > 0.8 minimal 极简样式-低调细线
|
||||
negative <= 0.8 card 卡片样式-半透明背景
|
||||
neutral 任意 default 默认样式-白字黑边
|
||||
|
||||
3. 使用ASS字幕格式支持每个字幕不同的样式
|
||||
4. FFmpeg将ASS字幕烧录到视频
|
||||
|
||||
使用方法:
|
||||
- 自动模式: 系统自动根据情感选择样式
|
||||
- 手动模式: 设置 force_style=True 强制使用指定样式
|
||||
""")
|
||||
|
||||
def main():
|
||||
"""主测试函数"""
|
||||
print_summary()
|
||||
|
||||
tests = [
|
||||
("情感分析", test_sentiment_analysis),
|
||||
("字幕生成", test_subtitle_generation),
|
||||
("ASS文件生成", test_ass_file_generation),
|
||||
]
|
||||
|
||||
passed = 0
|
||||
total = len(tests)
|
||||
|
||||
for test_name, test_func in tests:
|
||||
if test_func():
|
||||
passed += 1
|
||||
|
||||
print("\n" + "=" * 60)
|
||||
print(f"测试结果: {passed}/{total} 通过")
|
||||
print("=" * 60)
|
||||
|
||||
if passed == total:
|
||||
print("🎉 所有测试通过! 情感驱动字幕功能已就绪!")
|
||||
return 0
|
||||
else:
|
||||
print("⚠️ 部分测试未通过，请检查配置")
|
||||
return 1
|
||||
|
||||
if __name__ == '__main__':
|
||||
sys.exit(main())
|
||||
@@ -0,0 +1,943 @@
|
||||
#!/usr/bin/env python3
|
||||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
字幕封面生成器 - 后端处理逻辑
|
||||
"""
|
||||
|
||||
import os
|
||||
import sys
|
||||
import json
|
||||
import time
|
||||
import tempfile
|
||||
import subprocess
|
||||
from pathlib import Path
|
||||
from typing import Dict, Any, Optional, List, Tuple
|
||||
import logging
|
||||
|
||||
# ⚠️ 修复模块路径：确保 Python 能找到 modules 和其他本地模块
|
||||
# 将脚本所在目录（python 目录）添加到 sys.path
|
||||
script_dir = os.path.dirname(os.path.abspath(__file__))
|
||||
if script_dir not in sys.path:
|
||||
sys.path.insert(0, script_dir)
|
||||
print(f"[PATH] Added to sys.path: {script_dir}", file=sys.stderr)
|
||||
|
||||
try:
|
||||
import cv2
|
||||
import numpy as np
|
||||
from PIL import Image, ImageDraw, ImageFont, ImageFilter
|
||||
import jieba
|
||||
# import moviepy.editor # 暂时注释掉，需要时再启用
|
||||
except ImportError as e:
|
||||
print(f"缺少必要的依赖库: {e}")
|
||||
print("请安装以下库:")
|
||||
print("pip install opencv-python pillow jieba")
|
||||
sys.exit(1)
|
||||
|
||||
try:
|
||||
from modules.font_manager import get_font_manager
|
||||
FONT_MANAGER_AVAILABLE = True
|
||||
print("[PATH] font_manager loaded successfully", file=sys.stderr)
|
||||
except ImportError as e:
|
||||
FONT_MANAGER_AVAILABLE = False
|
||||
print(f"[PATH] Warning: font_manager module not available: {e}", file=sys.stderr)
|
||||
|
||||
# 配置日志
|
||||
logging.basicConfig(
|
||||
level=logging.INFO,
|
||||
format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
|
||||
)
|
||||
|
||||
class VideoSubtitleCoverGenerator:
|
||||
"""字幕封面生成器"""
|
||||
|
||||
def __init__(self, config: Dict[str, Any]):
|
||||
self.config = config
|
||||
self.logger = logging.getLogger(__name__)
|
||||
|
||||
# 字幕样式预设
|
||||
self.subtitle_presets = {
|
||||
"classic": {
|
||||
"fontFamily": "Arial",
|
||||
"fontSize": 24,
|
||||
"fontColor": "#FFFFFF",
|
||||
"backgroundColor": "#000000",
|
||||
"backgroundColorOpacity": 0.7,
|
||||
"position": "bottom"
|
||||
},
|
||||
"modern": {
|
||||
"fontFamily": "NotoSerifCJK-VF",
|
||||
"fontSize": 28,
|
||||
"fontColor": "#FFD700",
|
||||
"backgroundColor": "#1a1a1a",
|
||||
"backgroundColorOpacity": 0.8,
|
||||
"position": "bottom"
|
||||
},
|
||||
"minimal": {
|
||||
"fontFamily": "NotoSerifCJK-VF",
|
||||
"fontSize": 22,
|
||||
"fontColor": "#FFFFFF",
|
||||
"backgroundColor": "",
|
||||
"backgroundColorOpacity": 0,
|
||||
"position": "center"
|
||||
},
|
||||
"bold": {
|
||||
"fontFamily": "NotoSerifCJK-VF",
|
||||
"fontSize": 30,
|
||||
"fontColor": "#FF0000",
|
||||
"backgroundColor": "#FFFF00",
|
||||
"backgroundColorOpacity": 0.9,
|
||||
"position": "top"
|
||||
}
|
||||
}
|
||||
|
||||
# 封面样式预设
|
||||
self.cover_presets = {
|
||||
"big-text": {
|
||||
"timePointType": "random",
|
||||
"customTimePoint": 10,
|
||||
"textOverlay": "",
|
||||
"fontSize": 72,
|
||||
"fontColor": "white",
|
||||
"position": "center",
|
||||
"backgroundColor": True,
|
||||
"backgroundOpacity": 60,
|
||||
"maxWidth": 90,
|
||||
"effectType": "blur",
|
||||
"quality": 95
|
||||
},
|
||||
"search-card": {
|
||||
"timePointType": "middle",
|
||||
"customTimePoint": 50,
|
||||
"textOverlay": "",
|
||||
"fontSize": 48,
|
||||
"fontColor": "white",
|
||||
"position": "bottom-left",
|
||||
"backgroundColor": True,
|
||||
"backgroundOpacity": 80,
|
||||
"maxWidth": 70,
|
||||
"effectType": "none",
|
||||
"quality": 95
|
||||
},
|
||||
"cut-angle": {
|
||||
"timePointType": "start",
|
||||
"customTimePoint": 5,
|
||||
"textOverlay": "",
|
||||
"fontSize": 56,
|
||||
"fontColor": "yellow",
|
||||
"position": "bottom-right",
|
||||
"backgroundColor": False,
|
||||
"backgroundOpacity": 0,
|
||||
"maxWidth": 75,
|
||||
"effectType": "gradient",
|
||||
"quality": 95
|
||||
},
|
||||
"blur-person": {
|
||||
"timePointType": "random",
|
||||
"customTimePoint": 15,
|
||||
"textOverlay": "",
|
||||
"fontSize": 64,
|
||||
"fontColor": "white",
|
||||
"position": "bottom-center",
|
||||
"backgroundColor": True,
|
||||
"backgroundOpacity": 50,
|
||||
"maxWidth": 85,
|
||||
"effectType": "blur-person",
|
||||
"quality": 95
|
||||
}
|
||||
}
|
||||
|
||||
def get_subtitle_preset(self, preset_id: str) -> Dict[str, Any]:
|
||||
"""获取字幕预设样式"""
|
||||
return self.subtitle_presets.get(preset_id, self.subtitle_presets["classic"])
|
||||
|
||||
def get_cover_preset(self, preset_id: str) -> Dict[str, Any]:
|
||||
"""获取封面预设样式"""
|
||||
return self.cover_presets.get(preset_id, self.cover_presets["big-text"])
|
||||
|
||||
def extract_frame_at_time(self, video_path: str, time_seconds: float) -> Optional[np.ndarray]:
|
||||
"""从视频中提取指定时间的帧"""
|
||||
try:
|
||||
cap = cv2.VideoCapture(video_path)
|
||||
if not cap.isOpened():
|
||||
self.logger.error(f"无法打开视频文件: {video_path}")
|
||||
return None
|
||||
|
||||
# 设置时间位置
|
||||
cap.set(cv2.CAP_PROP_POS_MSEC, time_seconds * 1000)
|
||||
|
||||
ret, frame = cap.read()
|
||||
cap.release()
|
||||
|
||||
if ret:
|
||||
return frame
|
||||
else:
|
||||
self.logger.warning(f"无法读取视频帧: {time_seconds}s")
|
||||
return None
|
||||
|
||||
except Exception as e:
|
||||
self.logger.error(f"提取视频帧失败: {e}")
|
||||
return None
|
||||
|
||||
def get_video_duration(self, video_path: str) -> float:
|
||||
"""获取视频时长"""
|
||||
try:
|
||||
cap = cv2.VideoCapture(video_path)
|
||||
if not cap.isOpened():
|
||||
return 0
|
||||
|
||||
fps = cap.get(cv2.CAP_PROP_FPS)
|
||||
frame_count = cap.get(cv2.CAP_PROP_FRAME_COUNT)
|
||||
duration = frame_count / fps if fps > 0 else 0
|
||||
|
||||
cap.release()
|
||||
return duration
|
||||
|
||||
except Exception as e:
|
||||
self.logger.error(f"获取视频时长失败: {e}")
|
||||
return 0
|
||||
|
||||
def get_optimal_frame_time(self, video_path: str, time_point_type: str, custom_time: float) -> float:
|
||||
"""获取最佳帧时间点"""
|
||||
duration = self.get_video_duration(video_path)
|
||||
|
||||
if time_point_type == "start":
|
||||
return min(custom_time, duration * 0.1) # 前10%
|
||||
elif time_point_type == "middle":
|
||||
return duration * (custom_time / 100) # 百分比位置
|
||||
elif time_point_type == "end":
|
||||
return max(duration - custom_time, duration * 0.9) # 后10%
|
||||
elif time_point_type == "random":
|
||||
import random
|
||||
# 随机选择前60%的位置，避免片尾
|
||||
return random.uniform(0, duration * 0.6)
|
||||
else:
|
||||
return min(custom_time, duration * 0.5) # 默认中间位置
|
||||
|
||||
def apply_cover_effect(self, image: Image.Image, effect_type: str) -> Image.Image:
|
||||
"""应用封面特效"""
|
||||
if effect_type == "blur":
|
||||
# 高斯模糊背景
|
||||
return image.filter(ImageFilter.GaussianBlur(radius=5))
|
||||
elif effect_type == "blur-person":
|
||||
# 模拟人物虚化效果 (简单的中心区域保持清晰)
|
||||
width, height = image.size
|
||||
# 创建遮罩
|
||||
mask = Image.new('L', (width, height), 128)
|
||||
# 中心区域清晰
|
||||
center_x, center_y = width // 2, height // 2
|
||||
radius = min(width, height) // 4
|
||||
for x in range(width):
|
||||
for y in range(height):
|
||||
dist = ((x - center_x) ** 2 + (y - center_y) ** 2) ** 0.5
|
||||
if dist < radius:
|
||||
mask.putpixel((x, y), 255) # 清晰
|
||||
else:
|
||||
mask.putpixel((x, y), 64) # 模糊
|
||||
|
||||
blurred = image.filter(ImageFilter.GaussianBlur(radius=3))
|
||||
return Image.composite(image, blurred, mask)
|
||||
elif effect_type == "gradient":
|
||||
# 渐变效果
|
||||
width, height = image.size
|
||||
gradient = Image.new('L', (width, height), 0)
|
||||
draw = ImageDraw.Draw(gradient)
|
||||
|
||||
# 创建对角线渐变
|
||||
for x in range(width):
|
||||
for y in range(height):
|
||||
# 计算渐变值
|
||||
factor = (x + y) / (width + height)
|
||||
value = int(255 * (0.3 + 0.7 * factor)) # 30%-100% 不透明度
|
||||
gradient.putpixel((x, y), value)
|
||||
|
||||
# 创建半透明图层
|
||||
overlay = Image.new('RGBA', (width, height), (0, 0, 0, 0))
|
||||
gradient_rgba = Image.new('RGBA', (width, height), (0, 0, 0, 128))
|
||||
overlay = Image.composite(overlay, gradient_rgba, gradient)
|
||||
|
||||
return Image.alpha_composite(image.convert('RGBA'), overlay).convert('RGB')
|
||||
else:
|
||||
return image
|
||||
|
||||
def add_text_to_image(self, image: Image.Image, text: str, style_config: Dict[str, Any]) -> Image.Image:
|
||||
"""在图片上添加文字"""
|
||||
if not text.strip():
|
||||
return image
|
||||
|
||||
img = image.copy()
|
||||
draw = ImageDraw.Draw(img)
|
||||
width, height = img.size
|
||||
|
||||
# 获取字体（支持自定义字体）
|
||||
font_size = style_config.get('fontSize', 48)
|
||||
font_name = style_config.get('fontName', 'NotoSerifCJK-VF') # 从配置获取字体名称
|
||||
font_weight = style_config.get('fontWeight', 400)
|
||||
font = self._load_font_safe(font_name, font_size, font_weight)
|
||||
|
||||
# 计算文字尺寸
|
||||
bbox = draw.textbbox((0, 0), text, font=font)
|
||||
text_width = bbox[2] - bbox[0]
|
||||
text_height = bbox[3] - bbox[1]
|
||||
|
||||
# 限制最大宽度
|
||||
max_width_percent = style_config.get('maxWidth', 80) / 100
|
||||
max_width = int(width * max_width_percent)
|
||||
|
||||
if text_width > max_width:
|
||||
# 如果文字太宽，需要调整字体大小
|
||||
scale_factor = max_width / text_width
|
||||
new_font_size = int(font_size * scale_factor)
|
||||
try:
|
||||
font = ImageFont.truetype(font_family, new_font_size)
|
||||
except:
|
||||
font = ImageFont.load_default()
|
||||
|
||||
# 重新计算文字尺寸
|
||||
bbox = draw.textbbox((0, 0), text, font=font)
|
||||
text_width = bbox[2] - bbox[0]
|
||||
text_height = bbox[3] - bbox[1]
|
||||
|
||||
# 计算位置
|
||||
position = style_config.get('position', 'center')
|
||||
font_color = style_config.get('fontColor', 'white')
|
||||
|
||||
# 解析颜色
|
||||
if isinstance(font_color, str) and font_color.startswith('#'):
|
||||
r = int(font_color[1:3], 16)
|
||||
g = int(font_color[3:5], 16)
|
||||
b = int(font_color[5:7], 16)
|
||||
font_color = (r, g, b)
|
||||
elif font_color == 'white':
|
||||
font_color = (255, 255, 255)
|
||||
elif font_color == 'yellow':
|
||||
font_color = (255, 255, 0)
|
||||
else:
|
||||
font_color = (255, 255, 255)
|
||||
|
||||
# 计算文字位置
|
||||
if position == 'center':
|
||||
x = (width - text_width) // 2
|
||||
y = (height - text_height) // 2
|
||||
elif position == 'bottom-left':
|
||||
x = 50
|
||||
y = height - text_height - 50
|
||||
elif position == 'bottom-right':
|
||||
x = width - text_width - 50
|
||||
y = height - text_height - 50
|
||||
elif position == 'bottom-center':
|
||||
x = (width - text_width) // 2
|
||||
y = height - text_height - 50
|
||||
elif position == 'top-left':
|
||||
x = 50
|
||||
y = 50
|
||||
elif position == 'top-right':
|
||||
x = width - text_width - 50
|
||||
y = 50
|
||||
elif position == 'top-center':
|
||||
x = (width - text_width) // 2
|
||||
y = 50
|
||||
else: # 默认居中
|
||||
x = (width - text_width) // 2
|
||||
y = (height - text_height) // 2
|
||||
|
||||
# 添加背景
|
||||
background_color = style_config.get('backgroundColor', True)
|
||||
background_opacity = style_config.get('backgroundOpacity', 70)
|
||||
|
||||
if background_color and background_opacity > 0:
|
||||
# 创建半透明背景
|
||||
padding = 20
|
||||
bg_x1 = max(0, x - padding)
|
||||
bg_y1 = max(0, y - padding)
|
||||
bg_x2 = min(width, x + text_width + padding)
|
||||
bg_y2 = min(height, y + text_height + padding)
|
||||
|
||||
# 创建半透明黑色背景
|
||||
background = Image.new('RGBA', (width, height), (0, 0, 0, 0))
|
||||
bg_draw = ImageDraw.Draw(background)
|
||||
|
||||
# 计算背景颜色 (黑色半透明)
|
||||
alpha = int(255 * (background_opacity / 100))
|
||||
bg_color = (0, 0, 0, alpha)
|
||||
|
||||
bg_draw.rectangle([bg_x1, bg_y1, bg_x2, bg_y2], fill=bg_color)
|
||||
img = Image.alpha_composite(img.convert('RGBA'), background).convert('RGB')
|
||||
|
||||
# 重新创建draw对象
|
||||
draw = ImageDraw.Draw(img)
|
||||
|
||||
# 绘制文字
|
||||
draw.text((x, y), text, font=font, fill=font_color)
|
||||
|
||||
return img
|
||||
|
||||
def generate_cover_with_custom_template(self, video_path: str, output_path: str, template: Dict[str, Any]) -> bool:
|
||||
"""使用自定义模板生成封面"""
|
||||
try:
|
||||
# 提取视频的中间帧
|
||||
duration = self.get_video_duration(video_path)
|
||||
time_point = duration * 0.5 # 使用中间帧
|
||||
|
||||
frame = self.extract_frame_at_time(video_path, time_point)
|
||||
if frame is None:
|
||||
self.logger.error("无法提取视频帧")
|
||||
return False
|
||||
|
||||
# 转换为PIL图像
|
||||
image = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
|
||||
width, height = image.size
|
||||
|
||||
# 1. 应用背景模糊（如果启用）
|
||||
background_config = template.get('background', {})
|
||||
if background_config.get('blurEnabled', False):
|
||||
blur_radius = background_config.get('blurRadius', 15)
|
||||
image = image.filter(ImageFilter.GaussianBlur(radius=blur_radius))
|
||||
|
||||
# 2. 处理人像描边（如果启用）
|
||||
portrait_config = template.get('portrait', {})
|
||||
if portrait_config.get('strokeEnabled', False):
|
||||
# 创建绘图对象
|
||||
draw = ImageDraw.Draw(image)
|
||||
|
||||
# 获取人像位置和大小（百分比转换为像素）
|
||||
pos = portrait_config.get('position', {'x': 50, 'y': 50})
|
||||
size = portrait_config.get('size', {'width': 50, 'height': 70})
|
||||
|
||||
portrait_x = int(pos['x'] / 100 * width)
|
||||
portrait_y = int(pos['y'] / 100 * height)
|
||||
portrait_width = int(size['width'] / 100 * width)
|
||||
portrait_height = int(size['height'] / 100 * height)
|
||||
|
||||
# 绘制描边框
|
||||
stroke_color = portrait_config.get('strokeColor', '#FFD700')
|
||||
stroke_width = portrait_config.get('strokeWidth', 3)
|
||||
stroke_type = portrait_config.get('strokeType', 'solid')
|
||||
|
||||
# 计算矩形边界
|
||||
left = portrait_x - portrait_width // 2
|
||||
top = portrait_y - portrait_height // 2
|
||||
right = portrait_x + portrait_width // 2
|
||||
bottom = portrait_y + portrait_height // 2
|
||||
|
||||
if stroke_type == 'dashed':
|
||||
# 虚线描边
|
||||
dash_pattern = portrait_config.get('strokeDashPattern', [10, 5])
|
||||
# 简化实现：绘制虚线矩形
|
||||
# 这里需要手动绘制虚线，PIL不直接支持
|
||||
pass # 简化处理，暂时使用实线
|
||||
|
||||
# 绘制实线矩形
|
||||
for i in range(stroke_width):
|
||||
draw.rectangle(
|
||||
[left - i, top - i, right + i, bottom + i],
|
||||
outline=stroke_color,
|
||||
width=1
|
||||
)
|
||||
|
||||
# 3. 添加标题文字
|
||||
titles_config = template.get('titles', {})
|
||||
|
||||
# 主标题
|
||||
main_title = titles_config.get('main', {})
|
||||
main_text = main_title.get('text', '')
|
||||
if main_text:
|
||||
self._draw_text_with_stroke(
|
||||
image,
|
||||
main_text,
|
||||
main_title,
|
||||
width,
|
||||
height
|
||||
)
|
||||
|
||||
# 副标题
|
||||
sub_title = titles_config.get('sub', {})
|
||||
sub_text = sub_title.get('text', '')
|
||||
if sub_text:
|
||||
self._draw_text_with_stroke(
|
||||
image,
|
||||
sub_text,
|
||||
sub_title,
|
||||
width,
|
||||
height
|
||||
)
|
||||
|
||||
# 4. 添加蒙版（如果有）
|
||||
mask_config = template.get('mask', {})
|
||||
mask_path = mask_config.get('imagePath', '')
|
||||
if mask_path and os.path.exists(mask_path):
|
||||
try:
|
||||
mask_img = Image.open(mask_path).convert('RGBA')
|
||||
|
||||
# 获取蒙版位置和大小
|
||||
mask_pos = mask_config.get('position', {'x': 10, 'y': 10})
|
||||
mask_size = mask_config.get('size', {'width': 30, 'height': 20})
|
||||
mask_opacity = mask_config.get('opacity', 0.8)
|
||||
|
||||
mask_width = int(mask_size['width'] / 100 * width)
|
||||
mask_height = int(mask_size['height'] / 100 * height)
|
||||
mask_x = int(mask_pos['x'] / 100 * width)
|
||||
mask_y = int(mask_pos['y'] / 100 * height)
|
||||
|
||||
# 调整蒙版大小和透明度
|
||||
mask_img = mask_img.resize((mask_width, mask_height), Image.Resampling.LANCZOS)
|
||||
if mask_opacity < 1.0:
|
||||
alpha = mask_img.split()[3]
|
||||
alpha = alpha.point(lambda p: int(p * mask_opacity))
|
||||
mask_img.putalpha(alpha)
|
||||
|
||||
# 将蒙版粘贴到图像上
|
||||
image_rgba = image.convert('RGBA')
|
||||
image_rgba.paste(mask_img, (mask_x - mask_width // 2, mask_y - mask_height // 2), mask_img)
|
||||
image = image_rgba.convert('RGB')
|
||||
except Exception as e:
|
||||
self.logger.warning(f"添加蒙版失败: {e}")
|
||||
|
||||
# 保存图片
|
||||
image.save(output_path, 'JPEG', quality=95)
|
||||
self.logger.info(f"自定义模板封面生成成功: {output_path}")
|
||||
return True
|
||||
|
||||
except Exception as e:
|
||||
self.logger.error(f"生成自定义模板封面失败: {e}")
|
||||
return False
|
||||
|
||||
def _load_font_safe(self, font_name: str = None, font_size: int = 24, font_weight: int = 400) -> ImageFont.FreeTypeFont:
|
||||
"""
|
||||
加载字体，支持自定义字体、ziti目录、系统字体
|
||||
优先顺序：ziti目录 -> FontManager -> 系统字体 -> 默认字体
|
||||
|
||||
Args:
|
||||
font_name: 字体名称（如 "微软雅黑"、"墨趣古风体" 等）
|
||||
font_size: 字体大小
|
||||
font_weight: 字体粗细 (100-1000)
|
||||
|
||||
Returns:
|
||||
字体对象
|
||||
"""
|
||||
import platform
|
||||
|
||||
# 1. 首先直接尝试在 ziti 目录查找字体文件
|
||||
if font_name:
|
||||
# 多个可能的 ziti 目录路径（相对和绝对）
|
||||
possible_ziti_paths = [
|
||||
os.path.join(os.getcwd(), 'ziti'), # 当前工作目录
|
||||
os.path.join(os.path.dirname(__file__), '..', 'ziti'), # 相对于脚本
|
||||
os.path.join(os.path.dirname(__file__), '..', '..', 'ziti'), # 项目根目录
|
||||
'ziti', # 相对路径
|
||||
os.path.expanduser('~/ziti'), # 用户主目录
|
||||
]
|
||||
|
||||
# 尝试找到 ziti 目录
|
||||
for ziti_dir in possible_ziti_paths:
|
||||
if os.path.isdir(ziti_dir):
|
||||
self.logger.info(f"Found ziti directory: {ziti_dir}")
|
||||
|
||||
# 在 ziti 目录中查找字体文件
|
||||
for ext in ['.ttf', '.otf', '.ttc']:
|
||||
# 直接按字体名称查找
|
||||
font_path = os.path.join(ziti_dir, font_name + ext)
|
||||
if os.path.exists(font_path):
|
||||
try:
|
||||
return ImageFont.truetype(font_path, font_size)
|
||||
except Exception as e:
|
||||
self.logger.warning(f"Failed to load ziti font from {font_path}: {e}")
|
||||
|
||||
# 如果没有找到，扫描目录中所有文件进行模糊匹配
|
||||
try:
|
||||
for file in os.listdir(ziti_dir):
|
||||
if not file.startswith('.') and (file.endswith('.ttf') or file.endswith('.otf') or file.endswith('.ttc')):
|
||||
file_base = os.path.splitext(file)[0]
|
||||
# 检查是否匹配（精确或包含）
|
||||
if file_base == font_name or font_name in file_base:
|
||||
font_path = os.path.join(ziti_dir, file)
|
||||
try:
|
||||
return ImageFont.truetype(font_path, font_size)
|
||||
except Exception as e:
|
||||
self.logger.warning(f"Failed to load matched font {font_path}: {e}")
|
||||
except Exception as e:
|
||||
self.logger.warning(f"Error scanning ziti directory: {e}")
|
||||
|
||||
# 2. 如果 ziti 目录查找失败，使用 FontManager
|
||||
if font_name and FONT_MANAGER_AVAILABLE:
|
||||
try:
|
||||
font_manager = get_font_manager()
|
||||
font_path = font_manager.find_font(font_name, font_weight)
|
||||
|
||||
if font_path and os.path.exists(font_path):
|
||||
try:
|
||||
return ImageFont.truetype(font_path, font_size)
|
||||
except Exception as e:
|
||||
self.logger.warning(f"Failed to load custom font from {font_path}: {e}")
|
||||
except Exception as e:
|
||||
self.logger.warning(f"Font manager error: {e}")
|
||||
|
||||
# 3. 降级：根据操作系统加载默认系统字体
|
||||
system = platform.system()
|
||||
font_paths = []
|
||||
|
||||
if system == 'Windows':
|
||||
font_dir = None
|
||||
font_paths = [
|
||||
os.path.join(font_dir, 'NotoSerifCJK-VF'), # 微软雅黑
|
||||
'', # 微软雅黑 Bold
|
||||
'', # 黑体
|
||||
'', # 宋体
|
||||
]
|
||||
elif system == 'Darwin': # macOS
|
||||
font_paths = [
|
||||
'/Library/Fonts/PingFang.ttc',
|
||||
'/System/Library/Fonts/PingFang.ttc',
|
||||
]
|
||||
else: # Linux
|
||||
font_paths = [
|
||||
'/usr/share/fonts/truetype/wqy/wqy-zenhei.ttc',
|
||||
]
|
||||
|
||||
# 尝试加载每个字体
|
||||
for font_path in font_paths:
|
||||
if os.path.exists(font_path):
|
||||
try:
|
||||
return ImageFont.truetype(font_path, font_size)
|
||||
except:
|
||||
continue
|
||||
|
||||
return ImageFont.load_default()
|
||||
|
||||
def _draw_text_with_stroke(self, image: Image.Image, text: str, config: Dict[str, Any], img_width: int, img_height: int):
|
||||
"""在图片上绘制带描边的文字"""
|
||||
if not text:
|
||||
return
|
||||
|
||||
# 获取配置
|
||||
font_size = config.get('fontSize', 60)
|
||||
font_weight = config.get('fontWeight', 700)
|
||||
font_name = config.get('fontName', 'NotoSerifCJK-VF') # 从配置获取字体名称
|
||||
color = config.get('color', '#FFFFFF')
|
||||
stroke_color = config.get('strokeColor', '#000000')
|
||||
stroke_width = config.get('strokeWidth', 2)
|
||||
position = config.get('position', {'x': 50, 'y': 20})
|
||||
|
||||
# 加载字体（使用改进的字体加载逻辑，支持自定义字体）
|
||||
font = self._load_font_safe(font_name, font_size, font_weight)
|
||||
|
||||
# 计算文字位置
|
||||
text_x = int(position['x'] / 100 * img_width)
|
||||
text_y = int(position['y'] / 100 * img_height)
|
||||
|
||||
# 创建绘图对象
|
||||
draw = ImageDraw.Draw(image)
|
||||
|
||||
# 绘制描边（多次绘制以模拟描边效果）
|
||||
if stroke_width > 0:
|
||||
for adj_x in range(-stroke_width, stroke_width + 1):
|
||||
for adj_y in range(-stroke_width, stroke_width + 1):
|
||||
if adj_x != 0 or adj_y != 0:
|
||||
draw.text(
|
||||
(text_x + adj_x, text_y + adj_y),
|
||||
text,
|
||||
font=font,
|
||||
fill=stroke_color,
|
||||
anchor='mm'
|
||||
)
|
||||
|
||||
# 绘制主文字
|
||||
draw.text(
|
||||
(text_x, text_y),
|
||||
text,
|
||||
font=font,
|
||||
fill=color,
|
||||
anchor='mm'
|
||||
)
|
||||
|
||||
def generate_cover(self, video_path: str, output_path: str, cover_config: Dict[str, Any]) -> bool:
|
||||
"""生成封面图片"""
|
||||
try:
|
||||
# 获取最佳帧时间点
|
||||
time_point = self.get_optimal_frame_time(
|
||||
video_path,
|
||||
cover_config.get('timePointType', 'random'),
|
||||
cover_config.get('customTimePoint', 10)
|
||||
)
|
||||
|
||||
# 提取帧
|
||||
frame = self.extract_frame_at_time(video_path, time_point)
|
||||
if frame is None:
|
||||
self.logger.error("无法提取视频帧")
|
||||
return False
|
||||
|
||||
# 转换为PIL图像
|
||||
image = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
|
||||
|
||||
# 应用特效
|
||||
effect_type = cover_config.get('effectType', 'none')
|
||||
if effect_type != 'none':
|
||||
image = self.apply_cover_effect(image, effect_type)
|
||||
|
||||
# 添加文字
|
||||
text_overlay = cover_config.get('textOverlay', '').strip()
|
||||
if text_overlay:
|
||||
image = self.add_text_to_image(image, text_overlay, cover_config)
|
||||
|
||||
# 保存图片
|
||||
quality = cover_config.get('quality', 95)
|
||||
image.save(output_path, 'JPEG', quality=quality)
|
||||
|
||||
self.logger.info(f"封面生成成功: {output_path}")
|
||||
return True
|
||||
|
||||
except Exception as e:
|
||||
self.logger.error(f"生成封面失败: {e}")
|
||||
return False
|
||||
|
||||
def generate_subtitle_srt(self, text: str, video_duration: float) -> str:
|
||||
"""生成字幕SRT文件内容"""
|
||||
try:
|
||||
# 使用jieba分词
|
||||
words = list(jieba.cut(text.strip()))
|
||||
|
||||
# 计算字幕时长分布
|
||||
total_chars = len(text)
|
||||
if total_chars == 0:
|
||||
return ""
|
||||
|
||||
# 平均每个字幕显示2-3秒
|
||||
avg_duration = min(3.0, max(2.0, video_duration / max(1, total_chars / 10)))
|
||||
|
||||
srt_lines = []
|
||||
current_time = 1.0 # 从1秒开始
|
||||
|
||||
for i, word in enumerate(words):
|
||||
if not word.strip():
|
||||
continue
|
||||
|
||||
start_time = current_time
|
||||
end_time = min(current_time + avg_duration, video_duration - 0.5)
|
||||
|
||||
# 格式化时间 (HH:MM:SS,mmm)
|
||||
def format_time(seconds):
|
||||
hours = int(seconds // 3600)
|
||||
minutes = int((seconds % 3600) // 60)
|
||||
secs = int(seconds % 60)
|
||||
ms = int((seconds % 1) * 1000)
|
||||
return "02d"
|
||||
|
||||
start_str = format_time(start_time)
|
||||
end_str = format_time(end_time)
|
||||
|
||||
srt_lines.append(f"{i+1}")
|
||||
srt_lines.append(f"{start_str} --> {end_str}")
|
||||
srt_lines.append(word)
|
||||
srt_lines.append("") # 空行
|
||||
|
||||
current_time = end_time
|
||||
|
||||
return "\n".join(srt_lines)
|
||||
|
||||
except Exception as e:
|
||||
self.logger.error(f"生成字幕SRT失败: {e}")
|
||||
return ""
|
||||
|
||||
def add_subtitle_to_video(self, video_path: str, srt_path: str, output_path: str, subtitle_config: Dict[str, Any]) -> bool:
|
||||
"""为视频添加字幕 - 简化的实现"""
|
||||
try:
|
||||
# 目前先复制原视频作为输出（字幕功能可以后续完善）
|
||||
# 实际项目中可以使用 ffmpeg 或其他视频处理库
|
||||
import shutil
|
||||
shutil.copy2(video_path, output_path)
|
||||
|
||||
self.logger.info(f"视频复制成功 (字幕功能待完善): {output_path}")
|
||||
return True
|
||||
|
||||
except Exception as e:
|
||||
self.logger.error(f"复制视频失败: {e}")
|
||||
return False
|
||||
|
||||
def parse_srt_time(self, time_str: str) -> float:
|
||||
"""解析SRT时间格式"""
|
||||
try:
|
||||
# 格式: 00:00:01,000
|
||||
parts = time_str.replace(',', ':').split(':')
|
||||
if len(parts) >= 3:
|
||||
hours = int(parts[0])
|
||||
minutes = int(parts[1])
|
||||
seconds = float(parts[2])
|
||||
return hours * 3600 + minutes * 60 + seconds
|
||||
except:
|
||||
pass
|
||||
return 0.0
|
||||
|
||||
def process(self, progress_callback=None) -> Dict[str, Any]:
|
||||
"""主处理流程"""
|
||||
try:
|
||||
# 检查是否使用自定义模板
|
||||
if self.config.get('customTemplate', False):
|
||||
return self._process_with_custom_template(progress_callback)
|
||||
|
||||
model_config = self.config.get('modelConfig', {})
|
||||
video_path = model_config.get('video', '')
|
||||
|
||||
if not video_path or not os.path.exists(video_path):
|
||||
return {
|
||||
'code': -1,
|
||||
'msg': f'视频文件不存在: {video_path}'
|
||||
}
|
||||
|
||||
# 创建输出目录
|
||||
output_dir = os.path.dirname(video_path)
|
||||
base_name = os.path.splitext(os.path.basename(video_path))[0]
|
||||
|
||||
# 更新进度: 初始化
|
||||
if progress_callback:
|
||||
progress_callback(0.1, "初始化处理...")
|
||||
|
||||
# 获取字幕和封面配置
|
||||
subtitle_template = model_config.get('subtitleTemplate', 'classic')
|
||||
cover_template = model_config.get('coverTemplate', 'big-text')
|
||||
|
||||
subtitle_style = self.get_subtitle_preset(subtitle_template)
|
||||
cover_style = self.get_cover_preset(cover_template)
|
||||
|
||||
# 如果配置中有自定义样式，覆盖预设
|
||||
if 'subtitleStyle' in model_config:
|
||||
subtitle_style.update(model_config['subtitleStyle'])
|
||||
if 'coverStyle' in model_config:
|
||||
cover_style.update(model_config['coverStyle'])
|
||||
|
||||
# 更新进度: 生成字幕
|
||||
if progress_callback:
|
||||
progress_callback(0.2, "生成字幕文件...")
|
||||
|
||||
# 生成字幕SRT
|
||||
video_duration = self.get_video_duration(video_path)
|
||||
reference_text = model_config.get('referenceText', '')
|
||||
|
||||
srt_content = self.generate_subtitle_srt(reference_text, video_duration)
|
||||
|
||||
# 保存SRT文件
|
||||
srt_path = os.path.join(output_dir, f"{base_name}_subtitles.srt")
|
||||
with open(srt_path, 'w', encoding='utf-8') as f:
|
||||
f.write(srt_content)
|
||||
|
||||
# 更新进度: 生成封面
|
||||
if progress_callback:
|
||||
progress_callback(0.5, "生成视频封面...")
|
||||
|
||||
# 生成封面
|
||||
cover_path = os.path.join(output_dir, f"{base_name}_cover.jpg")
|
||||
cover_success = self.generate_cover(video_path, cover_path, cover_style)
|
||||
|
||||
# 更新进度: 添加字幕到视频
|
||||
if progress_callback:
|
||||
progress_callback(0.7, "为视频添加字幕...")
|
||||
|
||||
# 生成带字幕的视频
|
||||
video_with_subtitle_path = os.path.join(output_dir, f"{base_name}_with_subtitles.mp4")
|
||||
subtitle_success = self.add_subtitle_to_video(video_path, srt_path, video_with_subtitle_path, subtitle_style)
|
||||
|
||||
# 更新进度: 完成
|
||||
if progress_callback:
|
||||
progress_callback(1.0, "处理完成")
|
||||
|
||||
result = {
|
||||
'code': 0,
|
||||
'msg': '处理成功',
|
||||
'video': video_with_subtitle_path if subtitle_success else video_path,
|
||||
'cover': cover_path if cover_success else None,
|
||||
'srtPath': srt_path
|
||||
}
|
||||
|
||||
return result
|
||||
|
||||
except Exception as e:
|
||||
self.logger.error(f"处理失败: {e}")
|
||||
return {
|
||||
'code': -1,
|
||||
'msg': f'处理失败: {str(e)}'
|
||||
}
|
||||
|
||||
def _process_with_custom_template(self, progress_callback=None) -> Dict[str, Any]:
|
||||
"""使用自定义模板处理"""
|
||||
try:
|
||||
video_path = self.config.get('videoPath', '')
|
||||
template = self.config.get('template', {})
|
||||
|
||||
if not video_path or not os.path.exists(video_path):
|
||||
return {
|
||||
'code': -1,
|
||||
'msg': f'视频文件不存在: {video_path}'
|
||||
}
|
||||
|
||||
if not template:
|
||||
return {
|
||||
'code': -1,
|
||||
'msg': '缺少自定义模板配置'
|
||||
}
|
||||
|
||||
# 创建输出目录
|
||||
output_dir = self.config.get('outputDir', os.path.dirname(video_path))
|
||||
if not os.path.exists(output_dir):
|
||||
os.makedirs(output_dir, exist_ok=True)
|
||||
|
||||
base_name = os.path.splitext(os.path.basename(video_path))[0]
|
||||
|
||||
# 更新进度: 初始化
|
||||
if progress_callback:
|
||||
progress_callback(0.1, "初始化自定义模板处理...")
|
||||
|
||||
# 更新进度: 生成封面
|
||||
if progress_callback:
|
||||
progress_callback(0.5, "使用自定义模板生成封面...")
|
||||
|
||||
# 生成封面
|
||||
cover_path = os.path.join(output_dir, f"{base_name}_custom_cover.jpg")
|
||||
cover_success = self.generate_cover_with_custom_template(video_path, cover_path, template)
|
||||
|
||||
# 更新进度: 完成
|
||||
if progress_callback:
|
||||
progress_callback(1.0, "处理完成")
|
||||
|
||||
result = {
|
||||
'success': cover_success,
|
||||
'coverPath': cover_path if cover_success else None,
|
||||
'message': '自定义模板处理成功' if cover_success else '自定义模板处理失败'
|
||||
}
|
||||
|
||||
return result
|
||||
|
||||
except Exception as e:
|
||||
self.logger.error(f"自定义模板处理失败: {e}")
|
||||
return {
|
||||
'success': False,
|
||||
'error': str(e),
|
||||
'message': f'自定义模板处理失败: {str(e)}'
|
||||
}
|
||||
|
||||
|
||||
def main():
|
||||
"""命令行入口"""
|
||||
if len(sys.argv) < 2:
|
||||
print("用法: python video_subtitle_cover_generator.py <config.json>")
|
||||
sys.exit(1)
|
||||
|
||||
config_path = sys.argv[1]
|
||||
|
||||
try:
|
||||
with open(config_path, 'r', encoding='utf-8') as f:
|
||||
config = json.load(f)
|
||||
|
||||
generator = VideoSubtitleCoverGenerator(config)
|
||||
result = generator.process()
|
||||
|
||||
# 输出JSON结果（单行，以便Node.js解析）
|
||||
print(json.dumps(result, ensure_ascii=False))
|
||||
|
||||
except Exception as e:
|
||||
print(json.dumps({
|
||||
'success': False,
|
||||
'error': str(e),
|
||||
'message': f'执行失败: {str(e)}'
|
||||
}, ensure_ascii=False))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user