Skip to content
体验新版
项目
组织
正在加载...
登录
切换导航
打开侧边栏
灯光下的许诺
DeepMosaics
提交
84f6eb31
DeepMosaics
项目概览
灯光下的许诺
/
DeepMosaics
与 Fork 源项目一致
Fork自
Hypo / DeepMosaics
通知
7
Star
0
Fork
1
代码
文件
提交
分支
Tags
贡献者
分支图
Diff
Issue
0
列表
看板
标记
里程碑
合并请求
0
Wiki
0
Wiki
分析
仓库
DevOps
项目成员
Pages
DeepMosaics
项目概览
项目概览
详情
发布
仓库
仓库
文件
提交
分支
标签
贡献者
分支图
比较
Issue
0
Issue
0
列表
看板
标记
里程碑
合并请求
0
合并请求
0
Pages
分析
分析
仓库分析
DevOps
Wiki
0
Wiki
成员
成员
收起侧边栏
关闭侧边栏
动态
分支图
创建新Issue
提交
Issue看板
前往新版Gitcode,体验更适合开发者的 AI 搜索 >>
提交
84f6eb31
编写于
5月 10, 2020
作者:
H
hypox64
浏览文件
操作
浏览文件
下载
电子邮件补丁
差异文件
Create branch dev
上级
43f6674a
变更
16
隐藏空白更改
内联
并排
Showing
16 changed file
with
338 addition
and
314 deletion
+338
-314
.gitignore
.gitignore
+1
-0
README.md
README.md
+2
-2
README_CN.md
README_CN.md
+2
-2
cores/options.py
cores/options.py
+6
-4
deepmosaic.py
deepmosaic.py
+1
-1
docs/training_with_your_own_dataset.md
docs/training_with_your_own_dataset.md
+2
-2
make_datasets/make_pix2pix_dataset.py
make_datasets/make_pix2pix_dataset.py
+17
-45
make_datasets/make_video_dataset.py
make_datasets/make_video_dataset.py
+32
-49
models/loadmodel.py
models/loadmodel.py
+4
-22
models/pix2pixHD_model.py
models/pix2pixHD_model.py
+19
-1
models/runmodel.py
models/runmodel.py
+1
-1
train/add/train.py
train/add/train.py
+9
-10
train/clean/train.py
train/clean/train.py
+120
-88
util/data.py
util/data.py
+42
-10
util/image_processing.py
util/image_processing.py
+10
-37
util/mosaic.py
util/mosaic.py
+70
-40
未找到文件。
.gitignore
浏览文件 @
84f6eb31
...
...
@@ -141,6 +141,7 @@ dataset/
test*
video_tmp/
result/
nohup.out
#./
/pix2pix
/pix2pixHD
...
...
README.md
浏览文件 @
84f6eb31
...
...
@@ -67,11 +67,11 @@ You can download pre_trained models and put them into './pretrained_models'.<br>
#### Simple example
*
Add Mosaic (output media will save in './result')
<br>
```
bash
python3 deepmosaic.py
--media_path
./imgs/ruoruo.jpg
--model_path
./pretrained_models/mosaic/add_face.pth
--use_gpu
-1
python3 deepmosaic.py
--media_path
./imgs/ruoruo.jpg
--model_path
./pretrained_models/mosaic/add_face.pth
--use_gpu
0
```
*
Clean Mosaic (output media will save in './result')
<br>
```
bash
python3 deepmosaic.py
--media_path
./result/ruoruo_add.jpg
--model_path
./pretrained_models/mosaic/clean_face_HD.pth
--use_gpu
-1
python3 deepmosaic.py
--media_path
./result/ruoruo_add.jpg
--model_path
./pretrained_models/mosaic/clean_face_HD.pth
--use_gpu
0
```
#### More parameters
If you want to test other image or video, please refer to this file.
<br>
...
...
README_CN.md
浏览文件 @
84f6eb31
...
...
@@ -64,11 +64,11 @@ cd DeepMosaics
#### 简单的例子
*
为视频添加马赛克,例子中认为脸是需要打码的区域 ,可以通过切换预训练模型切换自动打码区域(输出结果将储存到 './result')
<br>
```
bash
python3 deepmosaic.py
--media_path
./imgs/ruoruo.jpg
--model_path
./pretrained_models/mosaic/add_face.pth
--use_gpu
-1
python3 deepmosaic.py
--media_path
./imgs/ruoruo.jpg
--model_path
./pretrained_models/mosaic/add_face.pth
--use_gpu
0
```
*
将视频中的马赛克移除,对于不同的打码物体需要使用对应的预训练模型进行马赛克消除(输出结果将储存到 './result')
<br>
```
bash
python3 deepmosaic.py
--media_path
./result/ruoruo_add.jpg
--model_path
./pretrained_models/mosaic/clean_face_HD.pth
--use_gpu
-1
python3 deepmosaic.py
--media_path
./result/ruoruo_add.jpg
--model_path
./pretrained_models/mosaic/clean_face_HD.pth
--use_gpu
0
```
#### 更多的参数
如果想要测试其他的图片或视频,请参照以下文件输入参数.
<br>
...
...
cores/options.py
浏览文件 @
84f6eb31
import
argparse
import
os
import
torch
class
Options
():
def
__init__
(
self
):
...
...
@@ -10,7 +10,7 @@ class Options():
def
initialize
(
self
):
#base
self
.
parser
.
add_argument
(
'--use_gpu'
,
type
=
int
,
default
=
0
,
help
=
'if -1,
do not use g
pu'
)
self
.
parser
.
add_argument
(
'--use_gpu'
,
type
=
int
,
default
=
0
,
help
=
'if -1,
use c
pu'
)
# self.parser.add_argument('--use_gpu', action='store_true', help='if input it, use gpu')
self
.
parser
.
add_argument
(
'--media_path'
,
type
=
str
,
default
=
'./imgs/ruoruo.jpg'
,
help
=
'your videos or images path'
)
self
.
parser
.
add_argument
(
'--mode'
,
type
=
str
,
default
=
'auto'
,
help
=
'Program running mode. auto | add | clean | style'
)
...
...
@@ -54,10 +54,12 @@ class Options():
model_name
=
os
.
path
.
basename
(
self
.
opt
.
model_path
)
os
.
environ
[
"CUDA_VISIBLE_DEVICES"
]
=
str
(
self
.
opt
.
use_gpu
)
import
torch
if
torch
.
cuda
.
is_available
()
and
self
.
opt
.
use_gpu
>
-
1
:
self
.
opt
.
use_gpu
=
True
pass
else
:
self
.
opt
.
use_gpu
=
False
self
.
opt
.
use_gpu
=
-
1
if
self
.
opt
.
mode
==
'auto'
:
if
'clean'
in
model_name
or
self
.
opt
.
traditional
:
...
...
deepmosaic.py
浏览文件 @
84f6eb31
...
...
@@ -72,6 +72,6 @@ if __name__ == '__main__':
for
stack
in
traceback
.
extract_tb
(
ex_stack
):
print
(
stack
)
input
(
'Please press any key to exit.
\n
'
)
util
.
clean_tempfiles
(
tmp_init
=
False
)
#
util.clean_tempfiles(tmp_init = False)
exit
(
0
)
docs/training_with_your_own_dataset.md
浏览文件 @
84f6eb31
...
...
@@ -53,7 +53,7 @@ python make_video_dataset.py --datadir 'dir for your videos' --model_path ../pre
### Add
```
bash
cd
train/add
python train.py
--
gpu_id
0
--dataset
../../datasets/draw/face
--savename
face
--loadsize
512
--finesize
360
--batchsize
16
python train.py
--
use_gpu
0
--dataset
../../datasets/draw/face
--savename
face
--loadsize
512
--finesize
360
--batchsize
16
```
### Clean
*
For image datasets(generated by
```make_pix2pix_dataset.py```
)
...
...
@@ -67,7 +67,7 @@ python train.py --name face --resize_or_crop resize_and_crop --loadSize 563 --fi
*
For video datasets(generated by
```make_video_dataset.py```
)
```
bash
cd
train/clean
python train.py
--dataset
../../datasets/video/face
--savename
face
--savefreq
100000
--gan
--hd
--lr
0.0002
--lambda_gan
1
--
gpu_id
0
--perload_num
8
python train.py
--dataset
../../datasets/video/face
--savename
face
--savefreq
100000
--gan
--hd
--lr
0.0002
--lambda_gan
1
--
use_gpu
0
```
## Testing
Put saved network to
```./pretrained_models/mosaic/```
and rename it as
```add_face.pth```
or
```clean_face_HD.pth```
or
```clean_face_video_HD.pth```
make_datasets/make_pix2pix_dataset.py
浏览文件 @
84f6eb31
import
os
import
random
import
sys
sys
.
path
.
append
(
".."
)
from
cores
import
Options
opt
=
Options
()
import
random
import
datetime
import
time
import
shutil
import
threading
import
warnings
warnings
.
filterwarnings
(
action
=
'ignore'
)
import
numpy
as
np
import
cv2
import
torch
sys
.
path
.
append
(
".."
)
from
models
import
runmodel
,
loadmodel
import
util.image_processing
as
impro
from
util
import
util
,
mosaic
,
data
from
cores
import
Options
opt
=
Options
()
opt
.
parser
.
add_argument
(
'--datadir'
,
type
=
str
,
default
=
'../datasets/draw/face'
,
help
=
''
)
opt
.
parser
.
add_argument
(
'--savedir'
,
type
=
str
,
default
=
'../datasets/pix2pix/face'
,
help
=
''
)
opt
.
parser
.
add_argument
(
'--name'
,
type
=
str
,
default
=
''
,
help
=
'save name'
)
...
...
@@ -61,51 +61,16 @@ if 'drawn' in opt.mod:
maskpaths
.
sort
()
if
'network'
in
opt
.
mod
or
'irregular'
in
opt
.
mod
:
imgpaths
=
util
.
Traversal
(
opt
.
datadir
)
imgpaths
=
util
.
is_imgs
(
imgpaths
)
random
.
shuffle
(
imgpaths
)
if
'irregular'
in
opt
.
mod
:
irrpaths
=
util
.
Traversal
(
opt
.
irrholedir
)
#def network
#def network
if
'network'
in
opt
.
mod
:
net
=
loadmodel
.
bisenet
(
opt
,
'roi'
)
# def checksaveimage(opt,img,mask):
# #check
# saveflag = True
# x,y,size,area = impro.boundingSquare(mask, random.uniform(1.4,1.6))
# if area < 1000:
# saveflag = False
# else:
# if opt.square:
# if size < opt.minsize:
# saveflag = False
# else:
# img = impro.resize(img[y-size:y+size,x-size:x+size],opt.outsize,interpolation=cv2.INTER_CUBIC)
# mask = impro.resize(mask[y-size:y+size,x-size:x+size],opt.outsize,interpolation=cv2.INTER_CUBIC)
# if impro.Q_lapulase(img)<opt.quality:
# saveflag = False
# else:
# img = impro.resize(img,opt.outsize,interpolation=cv2.INTER_CUBIC)
# if saveflag:
# # add mosaic
# img_mosaic = mosaic.addmosaic_random(img, mask)
# global savecnt
# savecnt += 1
# if opt.hd:
# cv2.imwrite(os.path.join(train_A_path,opt.name+'%06d' % savecnt+'.jpg'), img_mosaic)
# cv2.imwrite(os.path.join(train_B_path,opt.name+'%06d' % savecnt+'.jpg'), img)
# else:
# merge_img = impro.makedataset(img_mosaic, img)
# cv2.imwrite(os.path.join(train_path,opt.name+'%06d' % savecnt+'.jpg'), merge_img)
# if opt.savemask:
# cv2.imwrite(os.path.join(mask_save_path,opt.name+'%06d' % savecnt+'.png'), mask)
print
(
'Find images:'
,
len
(
imgpaths
))
starttime
=
datetime
.
datetime
.
now
()
filecnt
=
0
...
...
@@ -126,6 +91,8 @@ for fold in range(opt.fold):
mask
=
mask_irr
if
'network'
in
opt
.
mod
:
mask_net
=
runmodel
.
get_ROI_position
(
img
,
net
,
opt
,
keepsize
=
True
)[
0
]
if
opt
.
use_gpu
!=
-
1
:
torch
.
cuda
.
empty_cache
()
if
not
opt
.
all_mosaic_area
:
mask_net
=
impro
.
find_mostlikely_ROI
(
mask_net
)
mask
=
mask_net
...
...
@@ -157,7 +124,12 @@ for fold in range(opt.fold):
if
saveflag
:
# add mosaic
img_mosaic
=
mosaic
.
addmosaic_random
(
img
,
mask
)
# global savecnt
# random blur
if
random
.
random
()
>
0.5
:
Q
=
random
.
randint
(
1
,
15
)
img
=
impro
.
dctblur
(
img
,
Q
)
img_mosaic
=
impro
.
dctblur
(
img_mosaic
,
Q
)
savecnt
+=
1
if
opt
.
hd
:
...
...
@@ -185,6 +157,6 @@ for fold in range(opt.fold):
all_time
=
used_time
/
filecnt
*
all_length
print
(
'
\r
'
,
''
,
str
(
filecnt
)
+
'/'
+
str
(
all_length
)
+
' '
,
util
.
get_bar
(
percent
,
30
),
''
,
util
.
get_bar
(
percent
,
25
),
''
,
util
.
second2stamp
(
used_time
)
+
'/'
+
util
.
second2stamp
(
all_time
),
'f:'
+
str
(
savecnt
),
end
=
" "
)
\ No newline at end of file
make_datasets/make_video_dataset.py
浏览文件 @
84f6eb31
import
os
import
random
import
sys
sys
.
path
.
append
(
".."
)
from
cores
import
Options
opt
=
Options
()
import
random
import
datetime
import
time
import
shutil
import
threading
import
numpy
as
np
import
cv2
import
torch
sys
.
path
.
append
(
".."
)
from
models
import
runmodel
,
loadmodel
import
util.image_processing
as
impro
from
util
import
util
,
mosaic
,
data
,
ffmpeg
from
cores
import
Options
opt
=
Options
()
opt
.
parser
.
add_argument
(
'--datadir'
,
type
=
str
,
default
=
'your video dir'
,
help
=
''
)
opt
.
parser
.
add_argument
(
'--savedir'
,
type
=
str
,
default
=
'../datasets/video/face'
,
help
=
''
)
opt
.
parser
.
add_argument
(
'--interval'
,
type
=
int
,
default
=
30
,
help
=
'interval of split video '
)
...
...
@@ -25,6 +26,7 @@ opt.parser.add_argument('--quality', type=int ,default= 45,help='minimal quality
opt
.
parser
.
add_argument
(
'--outsize'
,
type
=
int
,
default
=
286
,
help
=
''
)
opt
.
parser
.
add_argument
(
'--startcnt'
,
type
=
int
,
default
=
0
,
help
=
''
)
opt
.
parser
.
add_argument
(
'--minsize'
,
type
=
int
,
default
=
96
,
help
=
'minimal roi size'
)
opt
.
parser
.
add_argument
(
'--no_sclectscene'
,
action
=
'store_true'
,
help
=
''
)
opt
=
opt
.
getparse
()
...
...
@@ -44,52 +46,31 @@ video_cnt = 1
starttime
=
datetime
.
datetime
.
now
()
for
videopath
in
videopaths
:
try
:
timestamps
=
[]
fps
,
endtime
,
height
,
width
=
ffmpeg
.
get_video_infos
(
videopath
)
for
cut_point
in
range
(
1
,
int
((
endtime
-
opt
.
time
)
/
opt
.
interval
)):
util
.
clean_tempfiles
()
ffmpeg
.
video2image
(
videopath
,
'./tmp/video2image/%05d.'
+
opt
.
tempimage_type
,
fps
=
1
,
start_time
=
util
.
second2stamp
(
cut_point
*
opt
.
interval
),
last_time
=
util
.
second2stamp
(
opt
.
time
))
imagepaths
=
util
.
Traversal
(
'./tmp/video2image'
)
cnt
=
0
for
i
in
range
(
opt
.
time
):
img
=
impro
.
imread
(
imagepaths
[
i
])
mask
=
runmodel
.
get_ROI_position
(
img
,
net
,
opt
,
keepsize
=
True
)[
0
]
if
not
opt
.
all_mosaic_area
:
mask
=
impro
.
find_mostlikely_ROI
(
mask
)
x
,
y
,
size
,
area
=
impro
.
boundingSquare
(
mask
,
Ex_mul
=
1
)
if
area
>
opt
.
minmaskarea
and
size
>
opt
.
minsize
and
impro
.
Q_lapulase
(
img
)
>
opt
.
quality
:
cnt
+=
1
if
cnt
==
opt
.
time
:
# print(second)
timestamps
.
append
(
util
.
second2stamp
(
cut_point
*
opt
.
interval
))
if
opt
.
no_sclectscene
:
timestamps
=
[
'00:00:00'
]
else
:
timestamps
=
[]
fps
,
endtime
,
height
,
width
=
ffmpeg
.
get_video_infos
(
videopath
)
for
cut_point
in
range
(
1
,
int
((
endtime
-
opt
.
time
)
/
opt
.
interval
)):
util
.
clean_tempfiles
()
ffmpeg
.
video2image
(
videopath
,
'./tmp/video2image/%05d.'
+
opt
.
tempimage_type
,
fps
=
1
,
start_time
=
util
.
second2stamp
(
cut_point
*
opt
.
interval
),
last_time
=
util
.
second2stamp
(
opt
.
time
))
imagepaths
=
util
.
Traversal
(
'./tmp/video2image'
)
cnt
=
0
for
i
in
range
(
opt
.
time
):
img
=
impro
.
imread
(
imagepaths
[
i
])
mask
=
runmodel
.
get_ROI_position
(
img
,
net
,
opt
,
keepsize
=
True
)[
0
]
if
not
opt
.
all_mosaic_area
:
mask
=
impro
.
find_mostlikely_ROI
(
mask
)
x
,
y
,
size
,
area
=
impro
.
boundingSquare
(
mask
,
Ex_mul
=
1
)
if
area
>
opt
.
minmaskarea
and
size
>
opt
.
minsize
and
impro
.
Q_lapulase
(
img
)
>
opt
.
quality
:
cnt
+=
1
if
cnt
==
opt
.
time
:
# print(second)
timestamps
.
append
(
util
.
second2stamp
(
cut_point
*
opt
.
interval
))
util
.
writelog
(
os
.
path
.
join
(
opt
.
savedir
,
'opt.txt'
),
videopath
+
'
\n
'
+
str
(
timestamps
))
#print(timestamps)
# util.clean_tempfiles()
# fps,endtime,height,width = ffmpeg.get_video_infos(videopath)
# # print(fps,endtime,height,width)
# ffmpeg.continuous_screenshot(videopath, './tmp/video2image', 1)
# # find where to cut
# print('Find where to cut...')
# timestamps=[]
# imagepaths = util.Traversal('./tmp/video2image')
# for second in range(int(endtime)):
# if second%opt.interval==0:
# cnt = 0
# for i in range(opt.time):
# img = impro.imread(imagepaths[second+i])
# mask = runmodel.get_ROI_position(img,net,opt)[0]
# if not opt.all_mosaic_area:
# mask = impro.find_mostlikely_ROI(mask)
# if impro.mask_area(mask) > opt.minmaskarea and impro.Q_lapulase(img)>opt.quality:
# # print(impro.mask_area(mask))
# cnt +=1
# if cnt == opt.time:
# # print(second)
# timestamps.append(util.second2stamp(second))
#generate datasets
print
(
'Generate datasets...'
)
for
timestamp
in
timestamps
:
...
...
@@ -143,3 +124,5 @@ for videopath in videopaths:
util
.
writelog
(
os
.
path
.
join
(
opt
.
savedir
,
'opt.txt'
),
videopath
+
'
\n
'
+
str
(
result_cnt
)
+
'
\n
'
+
str
(
e
))
video_cnt
+=
1
if
opt
.
use_gpu
!=
-
1
:
torch
.
cuda
.
empty_cache
()
models/loadmodel.py
浏览文件 @
84f6eb31
...
...
@@ -34,7 +34,7 @@ def pix2pix(opt):
show_paramsnumber
(
netG
,
'netG'
)
netG
.
load_state_dict
(
torch
.
load
(
opt
.
model_path
))
netG
.
eval
()
if
opt
.
use_gpu
:
if
opt
.
use_gpu
!=
-
1
:
netG
.
cuda
()
return
netG
...
...
@@ -60,7 +60,7 @@ def style(opt):
__patch_instance_norm_state_dict
(
state_dict
,
netG
,
key
.
split
(
'.'
))
netG
.
load_state_dict
(
state_dict
)
if
opt
.
use_gpu
:
if
opt
.
use_gpu
!=
-
1
:
netG
.
cuda
()
return
netG
...
...
@@ -72,7 +72,7 @@ def video(opt):
show_paramsnumber
(
netG
,
'netG'
)
netG
.
load_state_dict
(
torch
.
load
(
opt
.
model_path
))
netG
.
eval
()
if
opt
.
use_gpu
:
if
opt
.
use_gpu
!=
-
1
:
netG
.
cuda
()
return
netG
...
...
@@ -87,24 +87,6 @@ def bisenet(opt,type='roi'):
elif
type
==
'mosaic'
:
net
.
load_state_dict
(
torch
.
load
(
opt
.
mosaic_position_model_path
))
net
.
eval
()
if
opt
.
use_gpu
:
if
opt
.
use_gpu
!=
-
1
:
net
.
cuda
()
return
net
# def unet_clean(opt):
# net = UNet(n_channels = 3, n_classes = 1)
# show_paramsnumber(net,'segment')
# net.load_state_dict(torch.load(opt.mosaic_position_model_path))
# net.eval()
# if opt.use_gpu:
# net.cuda()
# return net
# def unet(opt):
# net = UNet(n_channels = 3, n_classes = 1)
# show_paramsnumber(net,'segment')
# net.load_state_dict(torch.load(opt.model_path))
# net.eval()
# if opt.use_gpu:
# net.cuda()
# return net
models/pix2pixHD_model.py
浏览文件 @
84f6eb31
...
...
@@ -48,7 +48,7 @@ def define_G(input_nc, output_nc, ngf, netG, n_downsample_global=3, n_blocks_glo
def
define_D
(
input_nc
,
ndf
,
n_layers_D
,
norm
=
'instance'
,
use_sigmoid
=
False
,
num_D
=
1
,
getIntermFeat
=
False
,
gpu_ids
=
[]):
norm_layer
=
get_norm_layer
(
norm_type
=
norm
)
netD
=
MultiscaleDiscriminator
(
input_nc
,
ndf
,
n_layers_D
,
norm_layer
,
use_sigmoid
,
num_D
,
getIntermFeat
)
print
(
netD
)
#
print(netD)
if
len
(
gpu_ids
)
>
0
:
assert
(
torch
.
cuda
.
is_available
())
netD
.
cuda
(
gpu_ids
[
0
])
...
...
@@ -67,6 +67,24 @@ def print_network(net):
##############################################################################
# Losses
##############################################################################
class
GAN_Feat_loss
(
nn
.
Module
):
def
__init__
(
self
,
opt
):
super
(
GAN_Feat_loss
,
self
).
__init__
()
self
.
num_D
=
opt
.
num_D
self
.
n_layers_D
=
opt
.
n_layers_D
self
.
lambda_feat
=
opt
.
lambda_feat
self
.
criterionFeat
=
nn
.
L1Loss
()
def
forward
(
self
,
pred_fake
,
pred_real
):
loss_G_GAN_Feat
=
0
feat_weights
=
4.0
/
(
self
.
n_layers_D
+
1
)
D_weights
=
1.0
/
self
.
num_D
for
i
in
range
(
self
.
num_D
):
for
j
in
range
(
len
(
pred_fake
[
i
])
-
1
):
loss_G_GAN_Feat
+=
D_weights
*
feat_weights
*
\
self
.
criterionFeat
(
pred_fake
[
i
][
j
],
pred_real
[
i
][
j
].
detach
())
*
self
.
lambda_feat
return
loss_G_GAN_Feat
class
GANLoss
(
nn
.
Module
):
def
__init__
(
self
,
use_lsgan
=
True
,
target_real_label
=
1.0
,
target_fake_label
=
0.0
,
tensor
=
torch
.
FloatTensor
):
...
...
models/runmodel.py
浏览文件 @
84f6eb31
...
...
@@ -7,7 +7,7 @@ from util import data
import
torch
import
numpy
as
np
def
run_segment
(
img
,
net
,
size
=
360
,
use_gpu
=
True
):
def
run_segment
(
img
,
net
,
size
=
360
,
use_gpu
=
0
):
img
=
impro
.
resize
(
img
,
size
)
img
=
data
.
im2tensor
(
img
,
use_gpu
=
use_gpu
,
bgr2rgb
=
False
,
use_transform
=
False
,
is0_1
=
True
)
mask
=
net
(
img
)
...
...
train/add/train.py
浏览文件 @
84f6eb31
import
sys
import
os
import
sys
sys
.
path
.
append
(
".."
)
sys
.
path
.
append
(
"../.."
)
from
cores
import
Options
opt
=
Options
()
import
random
import
datetime
import
time
...
...
@@ -13,9 +18,6 @@ import torch.backends.cudnn as cudnn
import
torch.nn
as
nn
from
torch
import
optim
sys
.
path
.
append
(
".."
)
sys
.
path
.
append
(
"../.."
)
from
cores
import
Options
from
util
import
mosaic
,
util
,
ffmpeg
,
filt
,
data
from
util
import
image_processing
as
impro
from
models
import
unet_model
,
BiSeNet_model
...
...
@@ -24,8 +26,6 @@ from models import unet_model,BiSeNet_model
'''
--------------------------Get options--------------------------
'''
opt
=
Options
()
opt
.
parser
.
add_argument
(
'--gpu_id'
,
type
=
int
,
default
=
0
,
help
=
''
)
opt
.
parser
.
add_argument
(
'--lr'
,
type
=
float
,
default
=
0.001
,
help
=
''
)
opt
.
parser
.
add_argument
(
'--finesize'
,
type
=
int
,
default
=
360
,
help
=
''
)
opt
.
parser
.
add_argument
(
'--loadsize'
,
type
=
int
,
default
=
400
,
help
=
''
)
...
...
@@ -51,12 +51,11 @@ dir_checkpoint = os.path.join('checkpoints/',opt.savename)
util
.
makedirs
(
dir_checkpoint
)
util
.
writelog
(
os
.
path
.
join
(
dir_checkpoint
,
'loss.txt'
),
str
(
time
.
asctime
(
time
.
localtime
(
time
.
time
())))
+
'
\n
'
+
util
.
opt2str
(
opt
))
torch
.
cuda
.
set_device
(
opt
.
gpu_id
)
def
Totensor
(
img
,
use_gpu
=
True
):
size
=
img
.
shape
[
0
]
img
=
torch
.
from_numpy
(
img
).
float
()
if
opt
.
use_gpu
:
if
opt
.
use_gpu
!=
-
1
:
img
=
img
.
cuda
()
return
img
...
...
@@ -110,7 +109,7 @@ if opt.continuetrain:
f
=
open
(
os
.
path
.
join
(
dir_checkpoint
,
'epoch_log.txt'
),
'r'
)
opt
.
startepoch
=
int
(
f
.
read
())
f
.
close
()
if
opt
.
use_gpu
:
if
opt
.
use_gpu
!=
-
1
:
net
.
cuda
()
cudnn
.
benchmark
=
True
...
...
@@ -134,7 +133,7 @@ for epoch in range(opt.startepoch,opt.maxepoch):
starttime
=
datetime
.
datetime
.
now
()
util
.
writelog
(
os
.
path
.
join
(
dir_checkpoint
,
'loss.txt'
),
'Epoch {}/{}.'
.
format
(
epoch
+
1
,
opt
.
maxepoch
),
True
)
net
.
train
()
if
opt
.
use_gpu
:
if
opt
.
use_gpu
!=
-
1
:
net
.
cuda
()
epoch_loss
=
0
for
i
in
range
(
int
(
img_num
*
0.8
/
opt
.
batchsize
)):
...
...
train/clean/train.py
浏览文件 @
84f6eb31
import
os
import
sys
sys
.
path
.
append
(
".."
)
sys
.
path
.
append
(
"../.."
)
from
cores
import
Options
opt
=
Options
()
import
numpy
as
np
import
cv2
import
random
...
...
@@ -6,12 +12,8 @@ import torch
import
torch.nn
as
nn
import
time
import
sys
sys
.
path
.
append
(
".."
)
sys
.
path
.
append
(
"../.."
)
from
util
import
mosaic
,
util
,
ffmpeg
,
filt
,
data
from
util
import
image_processing
as
impro
from
cores
import
Options
from
models
import
pix2pix_model
,
pix2pixHD_model
,
video_model
,
unet_model
,
loadmodel
,
videoHD_model
from
matplotlib
import
pyplot
as
plt
import
torch.backends.cudnn
as
cudnn
...
...
@@ -19,9 +21,6 @@ import torch.backends.cudnn as cudnn
'''
--------------------------Get options--------------------------
'''
opt
=
Options
()
opt
.
parser
.
add_argument
(
'--gpu_id'
,
type
=
int
,
default
=
0
,
help
=
''
)
opt
.
parser
.
add_argument
(
'--N'
,
type
=
int
,
default
=
25
,
help
=
''
)
opt
.
parser
.
add_argument
(
'--lr'
,
type
=
float
,
default
=
0.0002
,
help
=
''
)
opt
.
parser
.
add_argument
(
'--beta1'
,
type
=
float
,
default
=
0.5
,
help
=
''
)
...
...
@@ -33,8 +32,11 @@ opt.parser.add_argument('--lambda_gan',type=float,default=1, help='')
opt
.
parser
.
add_argument
(
'--finesize'
,
type
=
int
,
default
=
256
,
help
=
''
)
opt
.
parser
.
add_argument
(
'--loadsize'
,
type
=
int
,
default
=
286
,
help
=
''
)
opt
.
parser
.
add_argument
(
'--batchsize'
,
type
=
int
,
default
=
1
,
help
=
''
)
opt
.
parser
.
add_argument
(
'--perload_num'
,
type
=
int
,
default
=
16
,
help
=
'number of images pool'
)
opt
.
parser
.
add_argument
(
'--perload_num'
,
type
=
int
,
default
=
64
,
help
=
'number of images pool'
)
opt
.
parser
.
add_argument
(
'--norm'
,
type
=
str
,
default
=
'instance'
,
help
=
''
)
opt
.
parser
.
add_argument
(
'--num_D'
,
type
=
int
,
default
=
2
,
help
=
'number of discriminators to use'
)
opt
.
parser
.
add_argument
(
'--n_layers_D'
,
type
=
int
,
default
=
3
,
help
=
'only used if which_model_netD==n_layers'
)
opt
.
parser
.
add_argument
(
'--lambda_feat'
,
type
=
float
,
default
=
10.0
,
help
=
'weight for feature matching loss'
)
opt
.
parser
.
add_argument
(
'--dataset'
,
type
=
str
,
default
=
'./datasets/face/'
,
help
=
''
)
opt
.
parser
.
add_argument
(
'--maxiter'
,
type
=
int
,
default
=
10000000
,
help
=
''
)
...
...
@@ -52,11 +54,11 @@ dir_checkpoint = os.path.join('checkpoints/',opt.savename)
util
.
makedirs
(
dir_checkpoint
)
util
.
writelog
(
os
.
path
.
join
(
dir_checkpoint
,
'loss.txt'
),
str
(
time
.
asctime
(
time
.
localtime
(
time
.
time
())))
+
'
\n
'
+
util
.
opt2str
(
opt
))
torch
.
cuda
.
set_device
(
opt
.
gpu_id
)
cudnn
.
benchmark
=
True
N
=
opt
.
N
loss_sum
=
[
0.
,
0.
,
0.
,
0.
]
loss_plot
=
[[],[]]
loss_sum
=
[
0.
,
0.
,
0.
,
0.
,
0.
,
0
]
loss_plot
=
[[],[]
,[],[]
]
item_plot
=
[]
# list video dir
...
...
@@ -71,22 +73,24 @@ for video in videonames:
tmp
.
append
(
video
)
videonames
=
tmp
video_num
=
len
(
videonames
)
#def network
#--------------------------Init network--------------------------
print
(
'Init network...'
)
if
opt
.
hd
:
netG
=
videoHD_model
.
MosaicNet
(
3
*
N
+
1
,
3
,
norm
=
opt
.
norm
)
else
:
netG
=
video_model
.
MosaicNet
(
3
*
N
+
1
,
3
,
norm
=
opt
.
norm
)
netG
.
cuda
()
loadmodel
.
show_paramsnumber
(
netG
,
'netG'
)
if
opt
.
gan
:
if
opt
.
hd
:
#netD = pix2pixHD_model.define_D(6, 64, 3, norm = opt.norm, use_sigmoid=False, num_D=1)
netD
=
pix2pixHD_model
.
define_D
(
6
,
64
,
3
,
norm
=
opt
.
norm
,
use_sigmoid
=
False
,
num_D
=
2
,
getIntermFeat
=
True
)
netD
=
pix2pixHD_model
.
define_D
(
6
,
64
,
opt
.
n_layers_D
,
norm
=
opt
.
norm
,
use_sigmoid
=
False
,
num_D
=
opt
.
num_D
,
getIntermFeat
=
True
)
else
:
netD
=
pix2pix_model
.
define_D
(
3
*
2
,
64
,
'basic'
,
norm
=
opt
.
norm
)
netD
.
train
()
netD
.
cuda
()
#--------------------------continue train--------------------------
if
opt
.
continuetrain
:
if
not
os
.
path
.
isfile
(
os
.
path
.
join
(
dir_checkpoint
,
'last_G.pth'
)):
opt
.
continuetrain
=
False
...
...
@@ -99,58 +103,56 @@ if opt.continuetrain:
opt
.
startiter
=
int
(
f
.
read
())
f
.
close
()
#--------------------------optimizer & loss--------------------------
optimizer_G
=
torch
.
optim
.
Adam
(
netG
.
parameters
(),
lr
=
opt
.
lr
,
betas
=
(
opt
.
beta1
,
0.999
))
criterion_L1
=
nn
.
L1Loss
()
criterion_L2
=
nn
.
MSELoss
()
if
opt
.
gan
:
optimizer_D
=
torch
.
optim
.
Adam
(
netD
.
parameters
(),
lr
=
opt
.
lr
,
betas
=
(
opt
.
beta1
,
0.999
))
if
opt
.
hd
:
criterionGAN
=
pix2pixHD_model
.
GANLoss
(
tensor
=
torch
.
cuda
.
FloatTensor
)
criterionGAN
=
pix2pixHD_model
.
GANLoss
(
tensor
=
torch
.
cuda
.
FloatTensor
).
cuda
()
# criterionFeat = torch.nn.L1Loss()
criterionFeat
=
pix2pixHD_model
.
GAN_Feat_loss
(
opt
)
criterionVGG
=
pix2pixHD_model
.
VGGLoss
([
opt
.
use_gpu
])
else
:
criterionGAN
=
pix2pix_model
.
GANLoss
(
gan_mode
=
'lsgan'
).
cuda
()
if
opt
.
use_gpu
:
netG
.
cuda
()
if
opt
.
gan
:
netD
.
cuda
()
criterionGAN
.
cuda
()
cudnn
.
benchmark
=
True
'''
--------------------------preload data & data pool--------------------------
'''
def
loaddata
(
video_index
):
#
def loaddata(video_index):
videoname
=
videonames
[
video_index
]
img_index
=
random
.
randint
(
int
(
N
/
2
)
+
1
,
lengths
[
video_index
]
-
int
(
N
/
2
)
-
1
)
#
videoname = videonames[video_index]
#
img_index = random.randint(int(N/2)+1,lengths[video_index]- int(N/2)-1)
input_img
=
np
.
zeros
((
opt
.
loadsize
,
opt
.
loadsize
,
3
*
N
+
1
),
dtype
=
'uint8'
)
# this frame
this_mask
=
impro
.
imread
(
os
.
path
.
join
(
opt
.
dataset
,
videoname
,
'mask'
,
'%05d'
%
(
img_index
)
+
'.png'
),
'gray'
,
loadsize
=
opt
.
loadsize
)
input_img
[:,:,
-
1
]
=
this_mask
#print(os.path.join(opt.dataset,videoname,'origin_image','%05d'%(img_index)+'.jpg'))
ground_true
=
impro
.
imread
(
os
.
path
.
join
(
opt
.
dataset
,
videoname
,
'origin_image'
,
'%05d'
%
(
img_index
)
+
'.jpg'
),
loadsize
=
opt
.
loadsize
)
mosaic_size
,
mod
,
rect_rat
,
father
=
mosaic
.
get_random_parameter
(
ground_true
,
this_mask
)
# merge other frame
for
i
in
range
(
0
,
N
):
img
=
impro
.
imread
(
os
.
path
.
join
(
opt
.
dataset
,
videoname
,
'origin_image'
,
'%05d'
%
(
img_index
+
i
-
int
(
N
/
2
))
+
'.jpg'
),
loadsize
=
opt
.
loadsize
)
mask
=
impro
.
imread
(
os
.
path
.
join
(
opt
.
dataset
,
videoname
,
'mask'
,
'%05d'
%
(
img_index
+
i
-
int
(
N
/
2
))
+
'.png'
),
'gray'
,
loadsize
=
opt
.
loadsize
)
img_mosaic
=
mosaic
.
addmosaic_base
(
img
,
mask
,
mosaic_size
,
model
=
mod
,
rect_rat
=
rect_rat
,
father
=
father
)
input_img
[:,:,
i
*
3
:(
i
+
1
)
*
3
]
=
img_mosaic
# to tensor
input_img
,
ground_true
=
data
.
random_transform_video
(
input_img
,
ground_true
,
opt
.
finesize
,
N
)
input_img
=
data
.
im2tensor
(
input_img
,
bgr2rgb
=
False
,
use_gpu
=
opt
.
use_gpu
,
use_transform
=
False
,
is0_1
=
False
)
ground_true
=
data
.
im2tensor
(
ground_true
,
bgr2rgb
=
False
,
use_gpu
=
opt
.
use_gpu
,
use_transform
=
False
,
is0_1
=
False
)
# input_img = np.zeros((opt.loadsize,opt.loadsize,3*N+1), dtype='uint8')
# # this frame
# this_mask = impro.imread(os.path.join(opt.dataset,videoname,'mask','%05d'%(img_index)+'.png'),'gray',loadsize=opt.loadsize)
# input_img[:,:,-1] = this_mask
# #print(os.path.join(opt.dataset,videoname,'origin_image','%05d'%(img_index)+'.jpg'))
# ground_true = impro.imread(os.path.join(opt.dataset,videoname,'origin_image','%05d'%(img_index)+'.jpg'),loadsize=opt.loadsize)
# mosaic_size,mod,rect_rat,feather = mosaic.get_random_parameter(ground_true,this_mask)
# start_pos = mosaic.get_random_startpos(num=N,bisa_p=0.3,bisa_max=mosaic_size,bisa_max_part=3)
# # merge other frame
# for i in range(0,N):
# img = impro.imread(os.path.join(opt.dataset,videoname,'origin_image','%05d'%(img_index+i-int(N/2))+'.jpg'),loadsize=opt.loadsize)
# mask = impro.imread(os.path.join(opt.dataset,videoname,'mask','%05d'%(img_index+i-int(N/2))+'.png'),'gray',loadsize=opt.loadsize)
# img_mosaic = mosaic.addmosaic_base(img, mask, mosaic_size,model = mod,rect_rat=rect_rat,feather=feather,start_point=start_pos[i])
# input_img[:,:,i*3:(i+1)*3] = img_mosaic
# # to tensor
# input_img,ground_true = data.random_transform_video(input_img,ground_true,opt.finesize,N)
# input_img = data.im2tensor(input_img,bgr2rgb=False,use_gpu=-1,use_transform = False,is0_1=False)
# ground_true = data.im2tensor(ground_true,bgr2rgb=False,use_gpu=-1,use_transform = False,is0_1=False)
return
input_img
,
ground_true
#
return input_img,ground_true
print
(
'Preloading data, please wait...'
)
if
opt
.
perload_num
<=
opt
.
batchsize
:
opt
.
perload_num
=
opt
.
batchsize
*
2
#data pool
input_imgs
=
torch
.
rand
(
opt
.
perload_num
,
N
*
3
+
1
,
opt
.
finesize
,
opt
.
finesize
)
.
cuda
()
ground_trues
=
torch
.
rand
(
opt
.
perload_num
,
3
,
opt
.
finesize
,
opt
.
finesize
)
.
cuda
()
input_imgs
=
torch
.
rand
(
opt
.
perload_num
,
N
*
3
+
1
,
opt
.
finesize
,
opt
.
finesize
)
ground_trues
=
torch
.
rand
(
opt
.
perload_num
,
3
,
opt
.
finesize
,
opt
.
finesize
)
load_cnt
=
0
def
preload
():
...
...
@@ -158,8 +160,10 @@ def preload():
while
1
:
try
:
video_index
=
random
.
randint
(
0
,
video_num
-
1
)
ran
=
random
.
randint
(
0
,
opt
.
perload_num
-
1
)
input_imgs
[
ran
],
ground_trues
[
ran
]
=
loaddata
(
video_index
)
videoname
=
videonames
[
video_index
]
img_index
=
random
.
randint
(
int
(
N
/
2
)
+
1
,
lengths
[
video_index
]
-
int
(
N
/
2
)
-
1
)
input_imgs
[
load_cnt
%
opt
.
perload_num
],
ground_trues
[
load_cnt
%
opt
.
perload_num
]
=
data
.
load_train_video
(
videoname
,
img_index
,
opt
)
# input_imgs[load_cnt%opt.perload_num],ground_trues[load_cnt%opt.perload_num] = loaddata(video_index)
load_cnt
+=
1
# time.sleep(0.1)
except
Exception
as
e
:
...
...
@@ -172,7 +176,8 @@ time_start=time.time()
while
load_cnt
<
opt
.
perload_num
:
time
.
sleep
(
0.1
)
time_end
=
time
.
time
()
print
(
'load speed:'
,
round
((
time_end
-
time_start
)
/
opt
.
perload_num
,
3
),
's/it'
)
util
.
writelog
(
os
.
path
.
join
(
dir_checkpoint
,
'loss.txt'
),
'load speed: '
+
str
(
round
((
time_end
-
time_start
)
/
(
opt
.
perload_num
),
3
))
+
' s/it'
,
True
)
'''
--------------------------train--------------------------
...
...
@@ -180,22 +185,24 @@ print('load speed:',round((time_end-time_start)/opt.perload_num,3),'s/it')
util
.
copyfile
(
'./train.py'
,
os
.
path
.
join
(
dir_checkpoint
,
'train.py'
))
util
.
copyfile
(
'../../models/videoHD_model.py'
,
os
.
path
.
join
(
dir_checkpoint
,
'model.py'
))
netG
.
train
()
netD
.
train
()
time_start
=
time
.
time
()
print
(
"Begin training..."
)
for
iter
in
range
(
opt
.
startiter
+
1
,
opt
.
maxiter
):
ran
=
random
.
randint
(
0
,
opt
.
perload_num
-
opt
.
batchsize
-
1
)
inputdata
=
input_imgs
[
ran
:
ran
+
opt
.
batchsize
].
clone
()
target
=
ground_trues
[
ran
:
ran
+
opt
.
batchsize
].
clone
()
ran
=
random
.
randint
(
0
,
opt
.
perload_num
-
opt
.
batchsize
)
inputdata
=
(
input_imgs
[
ran
:
ran
+
opt
.
batchsize
].
clone
()).
cuda
()
target
=
(
ground_trues
[
ran
:
ran
+
opt
.
batchsize
].
clone
()).
cuda
()
if
opt
.
gan
:
# compute fake images: G(A)
pred
=
netG
(
inputdata
)
# update D
real_A
=
inputdata
[:,
int
((
N
-
1
)
/
2
)
*
3
:(
int
((
N
-
1
)
/
2
)
+
1
)
*
3
,:,:]
# --------------------update D--------------------
pix2pix_model
.
set_requires_grad
(
netD
,
True
)
optimizer_D
.
zero_grad
()
# Fake
real_A
=
inputdata
[:,
int
((
N
-
1
)
/
2
)
*
3
:(
int
((
N
-
1
)
/
2
)
+
1
)
*
3
,:,:]
fake_AB
=
torch
.
cat
((
real_A
,
pred
),
1
)
pred_fake
=
netD
(
fake_AB
.
detach
())
loss_D_fake
=
criterionGAN
(
pred_fake
,
False
)
...
...
@@ -205,29 +212,51 @@ for iter in range(opt.startiter+1,opt.maxiter):
loss_D_real
=
criterionGAN
(
pred_real
,
True
)
# combine loss and calculate gradients
loss_D
=
(
loss_D_fake
+
loss_D_real
)
*
0.5
loss_sum
[
2
]
+=
loss_D_fake
.
item
()
loss_sum
[
3
]
+=
loss_D_real
.
item
()
loss_sum
[
4
]
+=
loss_D_fake
.
item
()
loss_sum
[
5
]
+=
loss_D_real
.
item
()
# udpate D's weights
loss_D
.
backward
()
optimizer_D
.
step
()
#
update G
#
--------------------update G--------------------
pix2pix_model
.
set_requires_grad
(
netD
,
False
)
optimizer_G
.
zero_grad
()
# First, G(A) should fake the discriminator
real_A
=
inputdata
[:,
int
((
N
-
1
)
/
2
)
*
3
:(
int
((
N
-
1
)
/
2
)
+
1
)
*
3
,:,:]
fake_AB
=
torch
.
cat
((
real_A
,
pred
),
1
)
pred_fake
=
netD
(
fake_AB
)
loss_G_GAN
=
criterionGAN
(
pred_fake
,
True
)
*
opt
.
lambda_gan
# Second, G(A) = B
# GAN feature matching loss
# if opt.hd:
# real_AB = torch.cat((real_A, target), 1)
# pred_real = netD(real_AB)
# loss_G_GAN_Feat=criterionFeat(pred_fake,pred_real)
# loss_G_GAN_Feat = 0
# feat_weights = 4.0 / (opt.n_layers_D + 1)
# D_weights = 1.0 / opt.num_D
# for i in range(opt.num_D):
# for j in range(len(pred_fake[i])-1):
# loss_G_GAN_Feat += D_weights * feat_weights * criterionFeat(pred_fake[i][j], pred_real[i][j].detach()) * opt.lambda_feat
# combine loss and calculate gradients
if
opt
.
l2
:
loss_G_L1
=
(
criterion_L1
(
pred
,
target
)
+
criterion_L2
(
pred
,
target
))
*
opt
.
lambda_L1
else
:
loss_G_L1
=
criterion_L1
(
pred
,
target
)
*
opt
.
lambda_L1
# combine loss and calculate gradients
loss_G
=
loss_G_GAN
+
loss_G_L1
if
opt
.
hd
:
real_AB
=
torch
.
cat
((
real_A
,
target
),
1
)
pred_real
=
netD
(
real_AB
)
loss_G_GAN_Feat
=
criterionFeat
(
pred_fake
,
pred_real
)
loss_VGG
=
criterionVGG
(
pred
,
target
)
*
opt
.
lambda_feat
loss_G
=
loss_G_GAN
+
loss_G_L1
+
loss_G_GAN_Feat
+
loss_VGG
else
:
loss_G
=
loss_G_GAN
+
loss_G_L1
loss_sum
[
0
]
+=
loss_G_L1
.
item
()
loss_sum
[
1
]
+=
loss_G_GAN
.
item
()
loss_sum
[
2
]
+=
loss_G_GAN_Feat
.
item
()
loss_sum
[
3
]
+=
loss_VGG
.
item
()
# udpate G's weights
loss_G
.
backward
()
optimizer_G
.
step
()
...
...
@@ -244,64 +273,67 @@ for iter in range(opt.startiter+1,opt.maxiter):
loss_G_L1
.
backward
()
optimizer_G
.
step
()
if
(
iter
+
1
)
%
100
==
0
:
# save eval result
if
(
iter
+
1
)
%
1000
==
0
:
video_index
=
random
.
randint
(
0
,
video_num
-
1
)
videoname
=
videonames
[
video_index
]
img_index
=
random
.
randint
(
int
(
N
/
2
)
+
1
,
lengths
[
video_index
]
-
int
(
N
/
2
)
-
1
)
inputdata
,
target
=
data
.
load_train_video
(
videoname
,
img_index
,
opt
)
# inputdata,target = loaddata(random.randint(0,video_num-1))
inputdata
,
target
=
inputdata
.
cuda
(),
target
.
cuda
()
with
torch
.
no_grad
():
pred
=
netG
(
inputdata
)
try
:
data
.
showresult
(
inputdata
[:,
int
((
N
-
1
)
/
2
)
*
3
:(
int
((
N
-
1
)
/
2
)
+
1
)
*
3
,:,:],
target
,
pred
,
os
.
path
.
join
(
dir_checkpoint
,
'result_train
.jpg'
))
target
,
pred
,
os
.
path
.
join
(
dir_checkpoint
,
'result_eval
.jpg'
))
except
Exception
as
e
:
print
(
e
)
# plot
if
(
iter
+
1
)
%
1000
==
0
:
time_end
=
time
.
time
()
if
opt
.
gan
:
savestr
=
'iter:{0:d} L1_loss:{1:.
4f} G_loss:{2:.4f} D_f:{3:.4f} D_r:{4:.4
f} time:{5:.2f}'
.
format
(
savestr
=
'iter:{0:d} L1_loss:{1:.
3f} GAN_loss:{2:.3f} Feat:{3:.3f} VGG:{4:.3
f} time:{5:.2f}'
.
format
(
iter
+
1
,
loss_sum
[
0
]
/
1000
,
loss_sum
[
1
]
/
1000
,
loss_sum
[
2
]
/
1000
,
loss_sum
[
3
]
/
1000
,(
time_end
-
time_start
)
/
1000
)
util
.
writelog
(
os
.
path
.
join
(
dir_checkpoint
,
'loss.txt'
),
savestr
,
True
)
if
(
iter
+
1
)
/
1000
>=
10
:
loss_plot
[
0
].
append
(
loss_sum
[
0
]
/
1000
)
loss_plot
[
1
].
append
(
loss_sum
[
1
]
/
1000
)
item_plot
.
append
(
iter
+
1
)
try
:
plt
.
plot
(
item_plot
,
loss_plot
[
0
])
plt
.
plot
(
item_plot
,
loss_plot
[
1
])
plt
.
savefig
(
os
.
path
.
join
(
dir_checkpoint
,
'loss.jpg'
))
plt
.
close
()
except
Exception
as
e
:
print
(
"error:"
,
e
)
else
:
savestr
=
'iter:{0:d} L1_loss:{1:.4f} time:{2:.2f}'
.
format
(
iter
+
1
,
loss_sum
[
0
]
/
1000
,(
time_end
-
time_start
)
/
1000
)
util
.
writelog
(
os
.
path
.
join
(
dir_checkpoint
,
'loss.txt'
),
savestr
,
True
)
if
(
iter
+
1
)
/
1000
>=
10
:
loss_plot
[
0
].
append
(
loss_sum
[
0
]
/
1000
)
for
i
in
range
(
4
):
loss_plot
[
i
].
append
(
loss_sum
[
i
]
/
1000
)
item_plot
.
append
(
iter
+
1
)
try
:
plt
.
plot
(
item_plot
,
loss_plot
[
0
])
labels
=
[
'L1_loss'
,
'GAN_loss'
,
'GAN_Feat_loss'
,
'VGG_loss'
]
for
i
in
range
(
4
):
plt
.
plot
(
item_plot
,
loss_plot
[
i
],
label
=
labels
[
i
])
plt
.
xlabel
(
'iter'
)
plt
.
legend
(
loc
=
1
)
plt
.
savefig
(
os
.
path
.
join
(
dir_checkpoint
,
'loss.jpg'
))
plt
.
close
()
except
Exception
as
e
:
print
(
"error:"
,
e
)
loss_sum
=
[
0.
,
0.
,
0.
,
0.
]
loss_sum
=
[
0.
,
0.
,
0.
,
0.
,
0.
,
0.
]
time_start
=
time
.
time
()
# save network
if
(
iter
+
1
)
%
opt
.
savefreq
==
0
:
if
iter
+
1
!=
opt
.
savefreq
:
os
.
rename
(
os
.
path
.
join
(
dir_checkpoint
,
'last_G.pth'
),
os
.
path
.
join
(
dir_checkpoint
,
str
(
iter
+
1
-
opt
.
savefreq
)
+
'G.pth'
))
if
(
iter
+
1
)
%
(
opt
.
savefreq
//
10
)
==
0
:
torch
.
save
(
netG
.
cpu
().
state_dict
(),
os
.
path
.
join
(
dir_checkpoint
,
'last_G.pth'
))
if
opt
.
gan
:
if
iter
+
1
!=
opt
.
savefreq
:
os
.
rename
(
os
.
path
.
join
(
dir_checkpoint
,
'last_D.pth'
),
os
.
path
.
join
(
dir_checkpoint
,
str
(
iter
+
1
-
opt
.
savefreq
)
+
'D.pth'
))
torch
.
save
(
netD
.
cpu
().
state_dict
(),
os
.
path
.
join
(
dir_checkpoint
,
'last_D.pth'
))
if
opt
.
use_gpu
:
if
opt
.
use_gpu
!=-
1
:
netG
.
cuda
()
if
opt
.
gan
:
netD
.
cuda
()
f
=
open
(
os
.
path
.
join
(
dir_checkpoint
,
'iter'
),
'w+'
)
f
.
write
(
str
(
iter
+
1
))
f
.
close
()
if
(
iter
+
1
)
%
opt
.
savefreq
==
0
:
os
.
rename
(
os
.
path
.
join
(
dir_checkpoint
,
'last_G.pth'
),
os
.
path
.
join
(
dir_checkpoint
,
str
(
iter
+
1
)
+
'G.pth'
))
if
opt
.
gan
:
os
.
rename
(
os
.
path
.
join
(
dir_checkpoint
,
'last_D.pth'
),
os
.
path
.
join
(
dir_checkpoint
,
str
(
iter
+
1
)
+
'D.pth'
))
print
(
'network saved.'
)
#test
#test
if
(
iter
+
1
)
%
opt
.
savefreq
==
0
:
if
os
.
path
.
isdir
(
'./test'
):
netG
.
eval
()
...
...
util/data.py
浏览文件 @
84f6eb31
import
random
import
os
import
numpy
as
np
import
torch
import
torchvision.transforms
as
transforms
import
cv2
from
.
image_processing
import
color_adjust
,
dctblur
from
.
import
image_processing
as
impro
from
.
import
mosaic
transform
=
transforms
.
Compose
([
transforms
.
ToTensor
(),
transforms
.
Normalize
(
mean
=
(
0.5
,
0.5
,
0.5
),
std
=
(
0.5
,
0.5
,
0.5
))
...
...
@@ -34,7 +35,7 @@ def tensor2im(image_tensor, imtype=np.uint8, gray=False, rgb2bgr = True ,is0_1 =
return
image_numpy
.
astype
(
imtype
)
def
im2tensor
(
image_numpy
,
imtype
=
np
.
uint8
,
gray
=
False
,
bgr2rgb
=
True
,
reshape
=
True
,
use_gpu
=
True
,
use_transform
=
True
,
is0_1
=
True
):
def
im2tensor
(
image_numpy
,
imtype
=
np
.
uint8
,
gray
=
False
,
bgr2rgb
=
True
,
reshape
=
True
,
use_gpu
=
0
,
use_transform
=
True
,
is0_1
=
True
):
if
gray
:
h
,
w
=
image_numpy
.
shape
...
...
@@ -57,7 +58,7 @@ def im2tensor(image_numpy, imtype=np.uint8, gray=False,bgr2rgb = True, reshape =
image_tensor
=
torch
.
from_numpy
(
image_numpy
).
float
()
if
reshape
:
image_tensor
=
image_tensor
.
reshape
(
1
,
ch
,
h
,
w
)
if
use_gpu
:
if
use_gpu
!=
-
1
:
image_tensor
=
image_tensor
.
cuda
()
return
image_tensor
...
...
@@ -68,12 +69,18 @@ def shuffledata(data,target):
np
.
random
.
shuffle
(
target
)
def
random_transform_video
(
src
,
target
,
finesize
,
N
):
#random blur
if
random
.
random
()
<
0.2
:
h
,
w
=
src
.
shape
[:
2
]
src
=
src
[:
8
*
(
h
//
8
),:
8
*
(
w
//
8
)]
Q_ran
=
random
.
randint
(
1
,
15
)
src
[:,:,:
3
*
N
]
=
impro
.
dctblur
(
src
[:,:,:
3
*
N
],
Q_ran
)
target
=
impro
.
dctblur
(
target
,
Q_ran
)
#random crop
h
,
w
=
target
.
shape
[:
2
]
h_move
=
int
((
h
-
finesize
)
*
random
.
random
())
w_move
=
int
((
w
-
finesize
)
*
random
.
random
())
# print(h,w,h_move,w_move)
target
=
target
[
h_move
:
h_move
+
finesize
,
w_move
:
w_move
+
finesize
,:]
src
=
src
[
h_move
:
h_move
+
finesize
,
w_move
:
w_move
+
finesize
,:]
...
...
@@ -89,10 +96,10 @@ def random_transform_video(src,target,finesize,N):
g
=
random
.
uniform
(
-
0.05
,
0.05
)
r
=
random
.
uniform
(
-
0.05
,
0.05
)
for
i
in
range
(
N
):
src
[:,:,
i
*
3
:(
i
+
1
)
*
3
]
=
color_adjust
(
src
[:,:,
i
*
3
:(
i
+
1
)
*
3
],
alpha
,
beta
,
b
,
g
,
r
)
target
=
color_adjust
(
target
,
alpha
,
beta
,
b
,
g
,
r
)
src
[:,:,
i
*
3
:(
i
+
1
)
*
3
]
=
impro
.
color_adjust
(
src
[:,:,
i
*
3
:(
i
+
1
)
*
3
],
alpha
,
beta
,
b
,
g
,
r
)
target
=
impro
.
color_adjust
(
target
,
alpha
,
beta
,
b
,
g
,
r
)
#random blur
#random
resize
blur
if
random
.
random
()
<
0.5
:
interpolations
=
[
cv2
.
INTER_LINEAR
,
cv2
.
INTER_CUBIC
,
cv2
.
INTER_LANCZOS4
]
size_ran
=
random
.
uniform
(
0.7
,
1.5
)
...
...
@@ -156,7 +163,7 @@ def random_transform_image(img,mask,finesize,test_flag = False):
img
,
mask
=
img_crop
,
mask_crop
#random color
img
=
color_adjust
(
img
,
ran
=
True
)
img
=
impro
.
color_adjust
(
img
,
ran
=
True
)
#random flip
if
random
.
random
()
<
0.5
:
...
...
@@ -169,7 +176,7 @@ def random_transform_image(img,mask,finesize,test_flag = False):
#random blur
if
random
.
random
()
<
0.5
:
img
=
dctblur
(
img
,
random
.
randint
(
1
,
15
))
img
=
impro
.
dctblur
(
img
,
random
.
randint
(
1
,
15
))
# interpolations = [cv2.INTER_LINEAR,cv2.INTER_CUBIC,cv2.INTER_LANCZOS4]
# size_ran = random.uniform(0.7,1.5)
...
...
@@ -183,6 +190,31 @@ def random_transform_image(img,mask,finesize,test_flag = False):
print
(
'warning! shape error.'
)
return
img
,
mask
def
load_train_video
(
videoname
,
img_index
,
opt
):
N
=
opt
.
N
input_img
=
np
.
zeros
((
opt
.
loadsize
,
opt
.
loadsize
,
3
*
N
+
1
),
dtype
=
'uint8'
)
# this frame
this_mask
=
impro
.
imread
(
os
.
path
.
join
(
opt
.
dataset
,
videoname
,
'mask'
,
'%05d'
%
(
img_index
)
+
'.png'
),
'gray'
,
loadsize
=
opt
.
loadsize
)
input_img
[:,:,
-
1
]
=
this_mask
#print(os.path.join(opt.dataset,videoname,'origin_image','%05d'%(img_index)+'.jpg'))
ground_true
=
impro
.
imread
(
os
.
path
.
join
(
opt
.
dataset
,
videoname
,
'origin_image'
,
'%05d'
%
(
img_index
)
+
'.jpg'
),
loadsize
=
opt
.
loadsize
)
mosaic_size
,
mod
,
rect_rat
,
feather
=
mosaic
.
get_random_parameter
(
ground_true
,
this_mask
)
start_pos
=
mosaic
.
get_random_startpos
(
num
=
N
,
bisa_p
=
0.3
,
bisa_max
=
mosaic_size
,
bisa_max_part
=
3
)
# merge other frame
for
i
in
range
(
0
,
N
):
img
=
impro
.
imread
(
os
.
path
.
join
(
opt
.
dataset
,
videoname
,
'origin_image'
,
'%05d'
%
(
img_index
+
i
-
int
(
N
/
2
))
+
'.jpg'
),
loadsize
=
opt
.
loadsize
)
mask
=
impro
.
imread
(
os
.
path
.
join
(
opt
.
dataset
,
videoname
,
'mask'
,
'%05d'
%
(
img_index
+
i
-
int
(
N
/
2
))
+
'.png'
),
'gray'
,
loadsize
=
opt
.
loadsize
)
img_mosaic
=
mosaic
.
addmosaic_base
(
img
,
mask
,
mosaic_size
,
model
=
mod
,
rect_rat
=
rect_rat
,
feather
=
feather
,
start_point
=
start_pos
[
i
])
input_img
[:,:,
i
*
3
:(
i
+
1
)
*
3
]
=
img_mosaic
# to tensor
input_img
,
ground_true
=
random_transform_video
(
input_img
,
ground_true
,
opt
.
finesize
,
N
)
input_img
=
im2tensor
(
input_img
,
bgr2rgb
=
False
,
use_gpu
=-
1
,
use_transform
=
False
,
is0_1
=
False
)
ground_true
=
im2tensor
(
ground_true
,
bgr2rgb
=
False
,
use_gpu
=-
1
,
use_transform
=
False
,
is0_1
=
False
)
return
input_img
,
ground_true
def
showresult
(
img1
,
img2
,
img3
,
name
,
is0_1
=
False
):
size
=
img1
.
shape
[
3
]
showimg
=
np
.
zeros
((
size
,
size
*
3
,
3
))
...
...
util/image_processing.py
浏览文件 @
84f6eb31
...
...
@@ -117,44 +117,17 @@ def makedataset(target_image,orgin_image):
img
[
0
:
256
,
256
:
512
]
=
orgin_image
[
0
:
256
,
int
(
w
/
2
-
256
/
2
):
int
(
w
/
2
+
256
/
2
)]
return
img
def
spiltimage
(
img
,
size
=
128
):
h
,
w
=
img
.
shape
[:
2
]
# size = min(h,w)
if
w
>=
h
:
img1
=
img
[:,
0
:
size
]
img2
=
img
[:,
w
-
size
:
w
]
else
:
img1
=
img
[
0
:
size
,:]
img2
=
img
[
h
-
size
:
h
,:]
return
img1
,
img2
def
mergeimage
(
img1
,
img2
,
orgin_image
,
size
=
128
):
h
,
w
=
orgin_image
.
shape
[:
2
]
new_img1
=
np
.
zeros
((
h
,
w
),
dtype
=
"uint8"
)
new_img2
=
np
.
zeros
((
h
,
w
),
dtype
=
"uint8"
)
# size = min(h,w)
if
w
>=
h
:
new_img1
[:,
0
:
size
]
=
img1
new_img2
[:,
w
-
size
:
w
]
=
img2
else
:
new_img1
[
0
:
size
,:]
=
img1
new_img2
[
h
-
size
:
h
,:]
=
img2
result_img
=
cv2
.
add
(
new_img1
,
new_img2
)
return
result_img
def
block_dct_and_idct
(
g
,
QQF
):
T
=
cv2
.
dct
(
g
)
IT
=
np
.
round
(
cv2
.
idct
(
np
.
round
(
np
.
round
(
16.0
*
T
/
QQF
)
*
QQF
/
16
)))
return
IT
def
block_dct_and_idct
(
g
,
QQF
,
QQF_16
):
return
cv2
.
idct
(
np
.
round
(
16.0
*
cv2
.
dct
(
g
)
/
QQF
)
*
QQF_16
)
def
image_dct_and_idct
(
I
,
QF
):
h
,
w
=
I
.
shape
QQF
=
DCT_Q
*
QF
for
i
in
range
(
int
(
h
/
8
)):
for
j
in
range
(
int
(
w
/
8
)):
I
[
i
*
8
:(
i
+
1
)
*
8
,
j
*
8
:(
j
+
1
)
*
8
]
=
block_dct_and_idct
(
I
[
i
*
8
:(
i
+
1
)
*
8
,
j
*
8
:(
j
+
1
)
*
8
],
QQF
)
QQF_16
=
QQF
/
16.0
for
i
in
range
(
h
//
8
):
for
j
in
range
(
w
//
8
):
I
[
i
*
8
:(
i
+
1
)
*
8
,
j
*
8
:(
j
+
1
)
*
8
]
=
cv2
.
idct
(
np
.
round
(
16.0
*
cv2
.
dct
(
I
[
i
*
8
:(
i
+
1
)
*
8
,
j
*
8
:(
j
+
1
)
*
8
])
/
QQF
)
*
QQF_16
)
#I[i*8:(i+1)*8,j*8:(j+1)*8] = block_dct_and_idct(I[i*8:(i+1)*8,j*8:(j+1)*8],QQF,QQF_16)
return
I
def
dctblur
(
img
,
Q
):
...
...
@@ -162,7 +135,7 @@ def dctblur(img,Q):
Q: 1~20, 1->best
'''
h
,
w
=
img
.
shape
[:
2
]
img
[:
8
*
int
(
h
/
8
),:
8
*
int
(
w
/
8
)]
img
=
img
[:
8
*
(
h
//
8
),:
8
*
(
w
/
/
8
)]
img
=
img
.
astype
(
np
.
float32
)
if
img
.
ndim
==
2
:
img
=
image_dct_and_idct
(
img
,
Q
)
...
...
@@ -250,9 +223,9 @@ def Q_lapulase(resImg):
score
=
res
.
var
()
return
score
def
replace_mosaic
(
img_origin
,
img_fake
,
mask
,
x
,
y
,
size
,
no_father
):
def
replace_mosaic
(
img_origin
,
img_fake
,
mask
,
x
,
y
,
size
,
no_f
e
ather
):
img_fake
=
cv2
.
resize
(
img_fake
,(
size
*
2
,
size
*
2
),
interpolation
=
cv2
.
INTER_LANCZOS4
)
if
no_father
:
if
no_f
e
ather
:
img_origin
[
y
-
size
:
y
+
size
,
x
-
size
:
x
+
size
]
=
img_fake
img_result
=
img_origin
else
:
...
...
util/mosaic.py
浏览文件 @
84f6eb31
...
...
@@ -13,7 +13,7 @@ def addmosaic(img,mask,opt):
img
=
addmosaic_base
(
img
,
mask
,
opt
.
mosaic_size
,
opt
.
output_size
,
model
=
opt
.
mosaic_mod
)
return
img
def
addmosaic_base
(
img
,
mask
,
n
,
out_size
=
0
,
model
=
'squa_avg'
,
rect_rat
=
1.6
,
f
ather
=
0
):
def
addmosaic_base
(
img
,
mask
,
n
,
out_size
=
0
,
model
=
'squa_avg'
,
rect_rat
=
1.6
,
f
eather
=
0
,
start_point
=
[
0
,
0
]
):
'''
img: input image
mask: input mask
...
...
@@ -21,37 +21,50 @@ def addmosaic_base(img,mask,n,out_size = 0,model = 'squa_avg',rect_rat = 1.6,fat
out_size: output size 0->original
model : squa_avg squa_mid squa_random squa_avg_circle_edge rect_avg
rect_rat: if model==rect_avg , mosaic w/h=rect_rat
father : father size, -1->no 0->auto
feather : feather size, -1->no 0->auto
start_point : [0,0], please not input this parameter
'''
n
=
int
(
n
)
h_start
=
np
.
clip
(
start_point
[
0
],
0
,
n
)
w_start
=
np
.
clip
(
start_point
[
1
],
0
,
n
)
pix_mid_h
=
n
//
2
+
h_start
pix_mid_w
=
n
//
2
+
w_start
h
,
w
=
img
.
shape
[:
2
]
h_step
=
(
h
-
h_start
)
//
n
w_step
=
(
w
-
w_start
)
//
n
if
out_size
:
img
=
resize
(
img
,
out_size
)
h
,
w
=
img
.
shape
[:
2
]
mask
=
cv2
.
resize
(
mask
,(
w
,
h
))
img_mosaic
=
img
.
copy
()
if
mask
.
shape
[
0
]
!=
h
:
mask
=
cv2
.
resize
(
mask
,(
w
,
h
))
img_mosaic
=
img
.
copy
()
if
model
==
'squa_avg'
:
for
i
in
range
(
int
(
h
/
n
)):
for
j
in
range
(
int
(
w
/
n
)):
if
mask
[
int
(
i
*
n
+
n
/
2
),
int
(
j
*
n
+
n
/
2
)]
==
255
:
img_mosaic
[
i
*
n
:(
i
+
1
)
*
n
,
j
*
n
:(
j
+
1
)
*
n
,:]
=
img
[
i
*
n
:(
i
+
1
)
*
n
,
j
*
n
:(
j
+
1
)
*
n
,:].
mean
(
0
).
mean
(
0
)
for
i
in
range
(
h_step
):
for
j
in
range
(
w_step
):
if
mask
[
i
*
n
+
pix_mid_h
,
j
*
n
+
pix_mid_w
]:
img_mosaic
[
i
*
n
+
h_start
:(
i
+
1
)
*
n
+
h_start
,
j
*
n
+
w_start
:(
j
+
1
)
*
n
+
w_start
,:]
=
\
img
[
i
*
n
+
h_start
:(
i
+
1
)
*
n
+
h_start
,
j
*
n
+
w_start
:(
j
+
1
)
*
n
+
w_start
,:].
mean
(
axis
=
(
0
,
1
))
elif
model
==
'squa_mid'
:
for
i
in
range
(
int
(
h
/
n
)):
for
j
in
range
(
int
(
w
/
n
)):
if
mask
[
int
(
i
*
n
+
n
/
2
),
int
(
j
*
n
+
n
/
2
)]
==
255
:
img_mosaic
[
i
*
n
:(
i
+
1
)
*
n
,
j
*
n
:(
j
+
1
)
*
n
,:]
=
img
[
i
*
n
+
int
(
n
/
2
),
j
*
n
+
int
(
n
/
2
),:]
for
i
in
range
(
h_step
):
for
j
in
range
(
w_step
):
if
mask
[
i
*
n
+
pix_mid_h
,
j
*
n
+
pix_mid_w
]:
img_mosaic
[
i
*
n
+
h_start
:(
i
+
1
)
*
n
+
h_start
,
j
*
n
+
w_start
:(
j
+
1
)
*
n
+
w_start
,:]
=
\
img
[
i
*
n
+
n
//
2
+
h_start
,
j
*
n
+
n
//
2
+
w_start
,:]
elif
model
==
'squa_random'
:
for
i
in
range
(
int
(
h
/
n
)):
for
j
in
range
(
int
(
w
/
n
)):
if
mask
[
int
(
i
*
n
+
n
/
2
),
int
(
j
*
n
+
n
/
2
)]
==
255
:
img_mosaic
[
i
*
n
:(
i
+
1
)
*
n
,
j
*
n
:(
j
+
1
)
*
n
,:]
=
img
[
int
(
i
*
n
-
n
/
2
+
n
*
random
.
random
()),
int
(
j
*
n
-
n
/
2
+
n
*
random
.
random
()),:]
for
i
in
range
(
h_step
):
for
j
in
range
(
w_step
):
if
mask
[
i
*
n
+
pix_mid_h
,
j
*
n
+
pix_mid_w
]:
img_mosaic
[
i
*
n
+
h_start
:(
i
+
1
)
*
n
+
h_start
,
j
*
n
+
w_start
:(
j
+
1
)
*
n
+
w_start
,:]
=
\
img
[
h_start
+
int
(
i
*
n
-
n
/
2
+
n
*
random
.
random
()),
w_start
+
int
(
j
*
n
-
n
/
2
+
n
*
random
.
random
()),:]
elif
model
==
'squa_avg_circle_edge'
:
for
i
in
range
(
int
(
h
/
n
)):
for
j
in
range
(
int
(
w
/
n
)):
img_mosaic
[
i
*
n
:(
i
+
1
)
*
n
,
j
*
n
:(
j
+
1
)
*
n
,:]
=
img
[
i
*
n
:(
i
+
1
)
*
n
,
j
*
n
:(
j
+
1
)
*
n
,:].
mean
(
0
).
mean
(
0
)
for
i
in
range
(
h_step
):
for
j
in
range
(
w_step
):
img_mosaic
[
i
*
n
+
h_start
:(
i
+
1
)
*
n
+
h_start
,
j
*
n
+
w_start
:(
j
+
1
)
*
n
+
w_start
,:]
=
\
img
[
i
*
n
+
h_start
:(
i
+
1
)
*
n
+
h_start
,
j
*
n
+
w_start
:(
j
+
1
)
*
n
+
w_start
,:].
mean
(
axis
=
(
0
,
1
))
mask
=
cv2
.
threshold
(
mask
,
127
,
255
,
cv2
.
THRESH_BINARY
)[
1
]
_mask
=
ch_one2three
(
mask
)
mask_inv
=
cv2
.
bitwise_not
(
_mask
)
...
...
@@ -60,20 +73,24 @@ def addmosaic_base(img,mask,n,out_size = 0,model = 'squa_avg',rect_rat = 1.6,fat
img_mosaic
=
cv2
.
add
(
imgroi1
,
imgroi2
)
elif
model
==
'rect_avg'
:
n_h
=
n
n_w
=
int
(
n
*
rect_rat
)
for
i
in
range
(
int
(
h
/
n_h
)):
for
j
in
range
(
int
(
w
/
n_w
)):
if
mask
[
int
(
i
*
n_h
+
n_h
/
2
),
int
(
j
*
n_w
+
n_w
/
2
)]
==
255
:
img_mosaic
[
i
*
n_h
:(
i
+
1
)
*
n_h
,
j
*
n_w
:(
j
+
1
)
*
n_w
,:]
=
img
[
i
*
n_h
:(
i
+
1
)
*
n_h
,
j
*
n_w
:(
j
+
1
)
*
n_w
,:].
mean
(
0
).
mean
(
0
)
n_h
=
n
n_w
=
int
(
n
*
rect_rat
)
n_h_half
=
n_h
//
2
+
h_start
n_w_half
=
n_w
//
2
+
w_start
for
i
in
range
((
h
-
h_start
)
//
n_h
):
for
j
in
range
((
w
-
w_start
)
//
n_w
):
if
mask
[
i
*
n_h
+
n_h_half
,
j
*
n_w
+
n_w_half
]:
img_mosaic
[
i
*
n_h
+
h_start
:(
i
+
1
)
*
n_h
+
h_start
,
j
*
n_w
+
w_start
:(
j
+
1
)
*
n_w
+
w_start
,:]
=
\
img
[
i
*
n_h
+
h_start
:(
i
+
1
)
*
n_h
+
h_start
,
j
*
n_w
+
w_start
:(
j
+
1
)
*
n_w
+
w_start
,:].
mean
(
axis
=
(
0
,
1
))
if
father
!=
-
1
:
if
father
==
0
:
if
f
e
ather
!=
-
1
:
if
f
e
ather
==
0
:
mask
=
(
cv2
.
blur
(
mask
,
(
n
,
n
)))
else
:
mask
=
(
cv2
.
blur
(
mask
,
(
father
,
father
)))
mask
=
ch_one2three
(
mask
)
/
255.0
img_mosaic
=
(
img
*
(
1
-
mask
)
+
img_mosaic
*
mask
).
astype
(
'uint8'
)
mask
=
(
cv2
.
blur
(
mask
,
(
feather
,
feather
)))
mask
=
mask
/
255.0
for
i
in
range
(
3
):
img_mosaic
[:,:,
i
]
=
(
img
[:,:,
i
]
*
(
1
-
mask
)
+
img_mosaic
[:,:,
i
]
*
mask
)
img_mosaic
=
img_mosaic
.
astype
(
np
.
uint8
)
return
img_mosaic
...
...
@@ -108,19 +125,21 @@ def get_random_parameter(img,mask):
p
=
np
.
array
([
0.5
,
0.5
])
mod
=
np
.
random
.
choice
([
'normal'
,
'bounding'
],
p
=
p
.
ravel
())
mosaic_size
=
get_autosize
(
img
,
mask
,
area_type
=
mod
)
mosaic_size
=
int
(
mosaic_size
*
random
.
uniform
(
0.9
,
2.
1
))
mosaic_size
=
int
(
mosaic_size
*
random
.
uniform
(
0.9
,
2.
5
))
# mosaic mod
p
=
np
.
array
([
0.25
,
0.
25
,
0.1
,
0.4
])
mod
=
np
.
random
.
choice
([
'squa_mid'
,
'squa_avg'
,
'
squa_avg_circle_edge'
,
'
rect_avg'
],
p
=
p
.
ravel
())
p
=
np
.
array
([
0.25
,
0.
3
,
0.45
])
mod
=
np
.
random
.
choice
([
'squa_mid'
,
'squa_avg'
,
'rect_avg'
],
p
=
p
.
ravel
())
# rect_rat for rect_avg
rect_rat
=
random
.
uniform
(
1.1
,
1.6
)
# father size
father
=
int
(
mosaic_size
*
random
.
uniform
(
0
,
1.5
))
# feather size
feather
=
-
1
if
random
.
random
()
<
0.7
:
feather
=
int
(
mosaic_size
*
random
.
uniform
(
0
,
1.5
))
return
mosaic_size
,
mod
,
rect_rat
,
father
return
mosaic_size
,
mod
,
rect_rat
,
f
e
ather
def
addmosaic_autosize
(
img
,
mask
,
model
,
area_type
=
'normal'
):
...
...
@@ -129,6 +148,17 @@ def addmosaic_autosize(img,mask,model,area_type = 'normal'):
return
img_mosaic
def
addmosaic_random
(
img
,
mask
):
mosaic_size
,
mod
,
rect_rat
,
father
=
get_random_parameter
(
img
,
mask
)
img_mosaic
=
addmosaic_base
(
img
,
mask
,
mosaic_size
,
model
=
mod
,
rect_rat
=
rect_rat
,
father
=
father
)
return
img_mosaic
\ No newline at end of file
mosaic_size
,
mod
,
rect_rat
,
feather
=
get_random_parameter
(
img
,
mask
)
img_mosaic
=
addmosaic_base
(
img
,
mask
,
mosaic_size
,
model
=
mod
,
rect_rat
=
rect_rat
,
feather
=
feather
)
return
img_mosaic
def
get_random_startpos
(
num
,
bisa_p
,
bisa_max
,
bisa_max_part
):
pos
=
np
.
zeros
((
num
,
2
),
dtype
=
np
.
int64
)
if
random
.
random
()
<
bisa_p
:
indexs
=
random
.
sample
((
np
.
linspace
(
1
,
num
-
1
,
num
-
1
,
dtype
=
np
.
int64
)).
tolist
(),
random
.
randint
(
1
,
bisa_max_part
))
indexs
.
append
(
0
)
indexs
.
append
(
num
)
indexs
.
sort
()
for
i
in
range
(
len
(
indexs
)
-
1
):
pos
[
indexs
[
i
]:
indexs
[
i
+
1
]]
=
[
random
.
randint
(
0
,
bisa_max
),
random
.
randint
(
0
,
bisa_max
)]
return
pos
\ No newline at end of file
编辑
预览
Markdown
is supported
0%
请重试
或
添加新附件
.
添加附件
取消
You are about to add
0
people
to the discussion. Proceed with caution.
先完成此消息的编辑!
取消
想要评论请
注册
或
登录