Skip to content
体验新版
项目
组织
正在加载...
登录
切换导航
打开侧边栏
silvering
DeepMosaics
提交
cde29be4
DeepMosaics
项目概览
silvering
/
DeepMosaics
与 Fork 源项目一致
Fork自
Hypo / DeepMosaics
通知
1
Star
0
Fork
0
代码
文件
提交
分支
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,发现更多精彩内容 >>
提交
cde29be4
编写于
8月 14, 2019
作者:
H
hypox64
浏览文件
操作
浏览文件
下载
电子邮件补丁
差异文件
support for pix2pixHD model
上级
d790fbdb
变更
9
隐藏空白更改
内联
并排
Showing
9 changed file
with
467 addition
and
21 deletion
+467
-21
deepmosaic.py
deepmosaic.py
+5
-4
models/loadmodel.py
models/loadmodel.py
+6
-2
models/pix2pixHD_model.py
models/pix2pixHD_model.py
+418
-0
models/runmodel.py
models/runmodel.py
+6
-5
models/unet_parts.py
models/unet_parts.py
+1
-1
options.py
options.py
+8
-3
util/data.py
util/data.py
+2
-1
util/mosaic.py
util/mosaic.py
+13
-5
util/util.py
util/util.py
+8
-0
未找到文件。
deepmosaic.py
浏览文件 @
cde29be4
...
...
@@ -42,7 +42,7 @@ if opt.mode == 'add':
positions
.
append
([
x
,
y
,
area
])
cv2
.
imwrite
(
os
.
path
.
join
(
'./tmp/ROI_mask'
,
os
.
path
.
basename
(
imagepath
)),
mask
)
print
(
'Optimize
d
ROI locations...'
)
print
(
'Optimize ROI locations...'
)
mask_index
=
filt
.
position_medfilt
(
np
.
array
(
positions
),
7
)
# add mosaic
...
...
@@ -71,7 +71,7 @@ elif opt.mode == 'clean':
img_result
=
img_origin
.
copy
()
if
size
!=
0
:
img_mosaic
=
img_origin
[
y
-
size
:
y
+
size
,
x
-
size
:
x
+
size
]
img_fake
=
runmodel
.
run_pix2pix
(
img_mosaic
,
netG
,
use_gpu
=
opt
.
use_gpu
)
img_fake
=
runmodel
.
run_pix2pix
(
img_mosaic
,
netG
,
opt
)
img_result
=
impro
.
replace_mosaic
(
img_origin
,
img_fake
,
x
,
y
,
size
,
opt
.
no_feather
)
cv2
.
imwrite
(
os
.
path
.
join
(
opt
.
result_dir
,
os
.
path
.
basename
(
path
)),
img_result
)
...
...
@@ -90,7 +90,8 @@ elif opt.mode == 'clean':
img_origin
=
impro
.
imread
(
imagepath
)
x
,
y
,
size
=
runmodel
.
get_mosaic_position
(
img_origin
,
net_mosaic_pos
,
opt
)
positions
.
append
([
x
,
y
,
size
])
print
(
'Find Positions:'
,
imagepath
)
print
(
'Find mosaic location:'
,
imagepath
)
print
(
'Optimize mosaic locations...'
)
positions
=
np
.
array
(
positions
)
for
i
in
range
(
3
):
positions
[:,
i
]
=
filt
.
medfilt
(
positions
[:,
i
],
opt
.
medfilt_num
)
...
...
@@ -102,7 +103,7 @@ elif opt.mode == 'clean':
img_result
=
img_origin
.
copy
()
if
size
!=
0
:
img_mosaic
=
img_origin
[
y
-
size
:
y
+
size
,
x
-
size
:
x
+
size
]
img_fake
=
runmodel
.
run_pix2pix
(
img_mosaic
,
netG
,
use_gpu
=
opt
.
use_gpu
)
img_fake
=
runmodel
.
run_pix2pix
(
img_mosaic
,
netG
,
opt
)
img_result
=
impro
.
replace_mosaic
(
img_origin
,
img_fake
,
x
,
y
,
size
,
opt
.
no_feather
)
cv2
.
imwrite
(
os
.
path
.
join
(
'./tmp/replace_mosaic'
,
os
.
path
.
basename
(
imagepath
)),
img_result
)
print
(
'Clean Mosaic:'
,
imagepath
)
...
...
models/loadmodel.py
浏览文件 @
cde29be4
import
torch
from
.pix2pix_model
import
*
from
.pix2pix_model
import
define_G
from
.pix2pixHD_model
import
define_G
as
define_G_HD
from
.unet_model
import
UNet
def
pix2pix
(
opt
):
# print(opt.model_path,opt.netG)
netG
=
define_G
(
3
,
3
,
64
,
opt
.
netG
,
norm
=
'batch'
,
use_dropout
=
True
,
init_type
=
'normal'
,
gpu_ids
=
[])
if
opt
.
netG
==
'HD'
:
netG
=
define_G_HD
(
3
,
3
,
64
,
'global'
,
4
)
else
:
netG
=
define_G
(
3
,
3
,
64
,
opt
.
netG
,
norm
=
'batch'
,
use_dropout
=
True
,
init_type
=
'normal'
,
gpu_ids
=
[])
netG
.
load_state_dict
(
torch
.
load
(
opt
.
model_path
))
netG
.
eval
()
...
...
models/pix2pixHD_model.py
0 → 100644
浏览文件 @
cde29be4
# This code clone from https://github.com/NVIDIA/pix2pixHD
# LICENSE file : https://github.com/NVIDIA/pix2pixHD/blob/master/LICENSE.txt
import
torch
import
torch.nn
as
nn
import
functools
from
torch.autograd
import
Variable
import
numpy
as
np
###############################################################################
# Functions
###############################################################################
def
weights_init
(
m
):
classname
=
m
.
__class__
.
__name__
if
classname
.
find
(
'Conv'
)
!=
-
1
:
m
.
weight
.
data
.
normal_
(
0.0
,
0.02
)
elif
classname
.
find
(
'BatchNorm2d'
)
!=
-
1
:
m
.
weight
.
data
.
normal_
(
1.0
,
0.02
)
m
.
bias
.
data
.
fill_
(
0
)
def
get_norm_layer
(
norm_type
=
'instance'
):
if
norm_type
==
'batch'
:
norm_layer
=
functools
.
partial
(
nn
.
BatchNorm2d
,
affine
=
True
)
elif
norm_type
==
'instance'
:
norm_layer
=
functools
.
partial
(
nn
.
InstanceNorm2d
,
affine
=
False
)
else
:
raise
NotImplementedError
(
'normalization layer [%s] is not found'
%
norm_type
)
return
norm_layer
def
define_G
(
input_nc
,
output_nc
,
ngf
,
netG
,
n_downsample_global
=
3
,
n_blocks_global
=
9
,
n_local_enhancers
=
1
,
n_blocks_local
=
3
,
norm
=
'instance'
,
gpu_ids
=
[]):
norm_layer
=
get_norm_layer
(
norm_type
=
norm
)
if
netG
==
'global'
:
netG
=
GlobalGenerator
(
input_nc
,
output_nc
,
ngf
,
n_downsample_global
,
n_blocks_global
,
norm_layer
)
elif
netG
==
'local'
:
netG
=
LocalEnhancer
(
input_nc
,
output_nc
,
ngf
,
n_downsample_global
,
n_blocks_global
,
n_local_enhancers
,
n_blocks_local
,
norm_layer
)
elif
netG
==
'encoder'
:
netG
=
Encoder
(
input_nc
,
output_nc
,
ngf
,
n_downsample_global
,
norm_layer
)
else
:
raise
(
'generator not implemented!'
)
# print(netG)
if
len
(
gpu_ids
)
>
0
:
assert
(
torch
.
cuda
.
is_available
())
netG
.
cuda
(
gpu_ids
[
0
])
netG
.
apply
(
weights_init
)
return
netG
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
)
if
len
(
gpu_ids
)
>
0
:
assert
(
torch
.
cuda
.
is_available
())
netD
.
cuda
(
gpu_ids
[
0
])
netD
.
apply
(
weights_init
)
return
netD
def
print_network
(
net
):
if
isinstance
(
net
,
list
):
net
=
net
[
0
]
num_params
=
0
for
param
in
net
.
parameters
():
num_params
+=
param
.
numel
()
print
(
net
)
print
(
'Total number of parameters: %d'
%
num_params
)
##############################################################################
# Losses
##############################################################################
class
GANLoss
(
nn
.
Module
):
def
__init__
(
self
,
use_lsgan
=
True
,
target_real_label
=
1.0
,
target_fake_label
=
0.0
,
tensor
=
torch
.
FloatTensor
):
super
(
GANLoss
,
self
).
__init__
()
self
.
real_label
=
target_real_label
self
.
fake_label
=
target_fake_label
self
.
real_label_var
=
None
self
.
fake_label_var
=
None
self
.
Tensor
=
tensor
if
use_lsgan
:
self
.
loss
=
nn
.
MSELoss
()
else
:
self
.
loss
=
nn
.
BCELoss
()
def
get_target_tensor
(
self
,
input
,
target_is_real
):
target_tensor
=
None
if
target_is_real
:
create_label
=
((
self
.
real_label_var
is
None
)
or
(
self
.
real_label_var
.
numel
()
!=
input
.
numel
()))
if
create_label
:
real_tensor
=
self
.
Tensor
(
input
.
size
()).
fill_
(
self
.
real_label
)
self
.
real_label_var
=
Variable
(
real_tensor
,
requires_grad
=
False
)
target_tensor
=
self
.
real_label_var
else
:
create_label
=
((
self
.
fake_label_var
is
None
)
or
(
self
.
fake_label_var
.
numel
()
!=
input
.
numel
()))
if
create_label
:
fake_tensor
=
self
.
Tensor
(
input
.
size
()).
fill_
(
self
.
fake_label
)
self
.
fake_label_var
=
Variable
(
fake_tensor
,
requires_grad
=
False
)
target_tensor
=
self
.
fake_label_var
return
target_tensor
def
__call__
(
self
,
input
,
target_is_real
):
if
isinstance
(
input
[
0
],
list
):
loss
=
0
for
input_i
in
input
:
pred
=
input_i
[
-
1
]
target_tensor
=
self
.
get_target_tensor
(
pred
,
target_is_real
)
loss
+=
self
.
loss
(
pred
,
target_tensor
)
return
loss
else
:
target_tensor
=
self
.
get_target_tensor
(
input
[
-
1
],
target_is_real
)
return
self
.
loss
(
input
[
-
1
],
target_tensor
)
class
VGGLoss
(
nn
.
Module
):
def
__init__
(
self
,
gpu_ids
):
super
(
VGGLoss
,
self
).
__init__
()
self
.
vgg
=
Vgg19
().
cuda
()
self
.
criterion
=
nn
.
L1Loss
()
self
.
weights
=
[
1.0
/
32
,
1.0
/
16
,
1.0
/
8
,
1.0
/
4
,
1.0
]
def
forward
(
self
,
x
,
y
):
x_vgg
,
y_vgg
=
self
.
vgg
(
x
),
self
.
vgg
(
y
)
loss
=
0
for
i
in
range
(
len
(
x_vgg
)):
loss
+=
self
.
weights
[
i
]
*
self
.
criterion
(
x_vgg
[
i
],
y_vgg
[
i
].
detach
())
return
loss
##############################################################################
# Generator
##############################################################################
class
LocalEnhancer
(
nn
.
Module
):
def
__init__
(
self
,
input_nc
,
output_nc
,
ngf
=
32
,
n_downsample_global
=
3
,
n_blocks_global
=
9
,
n_local_enhancers
=
1
,
n_blocks_local
=
3
,
norm_layer
=
nn
.
BatchNorm2d
,
padding_type
=
'reflect'
):
super
(
LocalEnhancer
,
self
).
__init__
()
self
.
n_local_enhancers
=
n_local_enhancers
###### global generator model #####
ngf_global
=
ngf
*
(
2
**
n_local_enhancers
)
model_global
=
GlobalGenerator
(
input_nc
,
output_nc
,
ngf_global
,
n_downsample_global
,
n_blocks_global
,
norm_layer
).
model
model_global
=
[
model_global
[
i
]
for
i
in
range
(
len
(
model_global
)
-
3
)]
# get rid of final convolution layers
self
.
model
=
nn
.
Sequential
(
*
model_global
)
###### local enhancer layers #####
for
n
in
range
(
1
,
n_local_enhancers
+
1
):
### downsample
ngf_global
=
ngf
*
(
2
**
(
n_local_enhancers
-
n
))
model_downsample
=
[
nn
.
ReflectionPad2d
(
3
),
nn
.
Conv2d
(
input_nc
,
ngf_global
,
kernel_size
=
7
,
padding
=
0
),
norm_layer
(
ngf_global
),
nn
.
ReLU
(
True
),
nn
.
Conv2d
(
ngf_global
,
ngf_global
*
2
,
kernel_size
=
3
,
stride
=
2
,
padding
=
1
),
norm_layer
(
ngf_global
*
2
),
nn
.
ReLU
(
True
)]
### residual blocks
model_upsample
=
[]
for
i
in
range
(
n_blocks_local
):
model_upsample
+=
[
ResnetBlock
(
ngf_global
*
2
,
padding_type
=
padding_type
,
norm_layer
=
norm_layer
)]
### upsample
model_upsample
+=
[
nn
.
ConvTranspose2d
(
ngf_global
*
2
,
ngf_global
,
kernel_size
=
3
,
stride
=
2
,
padding
=
1
,
output_padding
=
1
),
norm_layer
(
ngf_global
),
nn
.
ReLU
(
True
)]
### final convolution
if
n
==
n_local_enhancers
:
model_upsample
+=
[
nn
.
ReflectionPad2d
(
3
),
nn
.
Conv2d
(
ngf
,
output_nc
,
kernel_size
=
7
,
padding
=
0
),
nn
.
Tanh
()]
setattr
(
self
,
'model'
+
str
(
n
)
+
'_1'
,
nn
.
Sequential
(
*
model_downsample
))
setattr
(
self
,
'model'
+
str
(
n
)
+
'_2'
,
nn
.
Sequential
(
*
model_upsample
))
self
.
downsample
=
nn
.
AvgPool2d
(
3
,
stride
=
2
,
padding
=
[
1
,
1
],
count_include_pad
=
False
)
def
forward
(
self
,
input
):
### create input pyramid
input_downsampled
=
[
input
]
for
i
in
range
(
self
.
n_local_enhancers
):
input_downsampled
.
append
(
self
.
downsample
(
input_downsampled
[
-
1
]))
### output at coarest level
output_prev
=
self
.
model
(
input_downsampled
[
-
1
])
### build up one layer at a time
for
n_local_enhancers
in
range
(
1
,
self
.
n_local_enhancers
+
1
):
model_downsample
=
getattr
(
self
,
'model'
+
str
(
n_local_enhancers
)
+
'_1'
)
model_upsample
=
getattr
(
self
,
'model'
+
str
(
n_local_enhancers
)
+
'_2'
)
input_i
=
input_downsampled
[
self
.
n_local_enhancers
-
n_local_enhancers
]
output_prev
=
model_upsample
(
model_downsample
(
input_i
)
+
output_prev
)
return
output_prev
class
GlobalGenerator
(
nn
.
Module
):
def
__init__
(
self
,
input_nc
,
output_nc
,
ngf
=
64
,
n_downsampling
=
3
,
n_blocks
=
9
,
norm_layer
=
nn
.
BatchNorm2d
,
padding_type
=
'reflect'
):
assert
(
n_blocks
>=
0
)
super
(
GlobalGenerator
,
self
).
__init__
()
activation
=
nn
.
ReLU
(
True
)
model
=
[
nn
.
ReflectionPad2d
(
3
),
nn
.
Conv2d
(
input_nc
,
ngf
,
kernel_size
=
7
,
padding
=
0
),
norm_layer
(
ngf
),
activation
]
### downsample
for
i
in
range
(
n_downsampling
):
mult
=
2
**
i
model
+=
[
nn
.
Conv2d
(
ngf
*
mult
,
ngf
*
mult
*
2
,
kernel_size
=
3
,
stride
=
2
,
padding
=
1
),
norm_layer
(
ngf
*
mult
*
2
),
activation
]
### resnet blocks
mult
=
2
**
n_downsampling
for
i
in
range
(
n_blocks
):
model
+=
[
ResnetBlock
(
ngf
*
mult
,
padding_type
=
padding_type
,
activation
=
activation
,
norm_layer
=
norm_layer
)]
### upsample
for
i
in
range
(
n_downsampling
):
mult
=
2
**
(
n_downsampling
-
i
)
model
+=
[
nn
.
ConvTranspose2d
(
ngf
*
mult
,
int
(
ngf
*
mult
/
2
),
kernel_size
=
3
,
stride
=
2
,
padding
=
1
,
output_padding
=
1
),
norm_layer
(
int
(
ngf
*
mult
/
2
)),
activation
]
model
+=
[
nn
.
ReflectionPad2d
(
3
),
nn
.
Conv2d
(
ngf
,
output_nc
,
kernel_size
=
7
,
padding
=
0
),
nn
.
Tanh
()]
self
.
model
=
nn
.
Sequential
(
*
model
)
def
forward
(
self
,
input
):
return
self
.
model
(
input
)
# Define a resnet block
class
ResnetBlock
(
nn
.
Module
):
def
__init__
(
self
,
dim
,
padding_type
,
norm_layer
,
activation
=
nn
.
ReLU
(
True
),
use_dropout
=
False
):
super
(
ResnetBlock
,
self
).
__init__
()
self
.
conv_block
=
self
.
build_conv_block
(
dim
,
padding_type
,
norm_layer
,
activation
,
use_dropout
)
def
build_conv_block
(
self
,
dim
,
padding_type
,
norm_layer
,
activation
,
use_dropout
):
conv_block
=
[]
p
=
0
if
padding_type
==
'reflect'
:
conv_block
+=
[
nn
.
ReflectionPad2d
(
1
)]
elif
padding_type
==
'replicate'
:
conv_block
+=
[
nn
.
ReplicationPad2d
(
1
)]
elif
padding_type
==
'zero'
:
p
=
1
else
:
raise
NotImplementedError
(
'padding [%s] is not implemented'
%
padding_type
)
conv_block
+=
[
nn
.
Conv2d
(
dim
,
dim
,
kernel_size
=
3
,
padding
=
p
),
norm_layer
(
dim
),
activation
]
if
use_dropout
:
conv_block
+=
[
nn
.
Dropout
(
0.5
)]
p
=
0
if
padding_type
==
'reflect'
:
conv_block
+=
[
nn
.
ReflectionPad2d
(
1
)]
elif
padding_type
==
'replicate'
:
conv_block
+=
[
nn
.
ReplicationPad2d
(
1
)]
elif
padding_type
==
'zero'
:
p
=
1
else
:
raise
NotImplementedError
(
'padding [%s] is not implemented'
%
padding_type
)
conv_block
+=
[
nn
.
Conv2d
(
dim
,
dim
,
kernel_size
=
3
,
padding
=
p
),
norm_layer
(
dim
)]
return
nn
.
Sequential
(
*
conv_block
)
def
forward
(
self
,
x
):
out
=
x
+
self
.
conv_block
(
x
)
return
out
class
Encoder
(
nn
.
Module
):
def
__init__
(
self
,
input_nc
,
output_nc
,
ngf
=
32
,
n_downsampling
=
4
,
norm_layer
=
nn
.
BatchNorm2d
):
super
(
Encoder
,
self
).
__init__
()
self
.
output_nc
=
output_nc
model
=
[
nn
.
ReflectionPad2d
(
3
),
nn
.
Conv2d
(
input_nc
,
ngf
,
kernel_size
=
7
,
padding
=
0
),
norm_layer
(
ngf
),
nn
.
ReLU
(
True
)]
### downsample
for
i
in
range
(
n_downsampling
):
mult
=
2
**
i
model
+=
[
nn
.
Conv2d
(
ngf
*
mult
,
ngf
*
mult
*
2
,
kernel_size
=
3
,
stride
=
2
,
padding
=
1
),
norm_layer
(
ngf
*
mult
*
2
),
nn
.
ReLU
(
True
)]
### upsample
for
i
in
range
(
n_downsampling
):
mult
=
2
**
(
n_downsampling
-
i
)
model
+=
[
nn
.
ConvTranspose2d
(
ngf
*
mult
,
int
(
ngf
*
mult
/
2
),
kernel_size
=
3
,
stride
=
2
,
padding
=
1
,
output_padding
=
1
),
norm_layer
(
int
(
ngf
*
mult
/
2
)),
nn
.
ReLU
(
True
)]
model
+=
[
nn
.
ReflectionPad2d
(
3
),
nn
.
Conv2d
(
ngf
,
output_nc
,
kernel_size
=
7
,
padding
=
0
),
nn
.
Tanh
()]
self
.
model
=
nn
.
Sequential
(
*
model
)
def
forward
(
self
,
input
,
inst
):
outputs
=
self
.
model
(
input
)
# instance-wise average pooling
outputs_mean
=
outputs
.
clone
()
inst_list
=
np
.
unique
(
inst
.
cpu
().
numpy
().
astype
(
int
))
for
i
in
inst_list
:
for
b
in
range
(
input
.
size
()[
0
]):
indices
=
(
inst
[
b
:
b
+
1
]
==
int
(
i
)).
nonzero
()
# n x 4
for
j
in
range
(
self
.
output_nc
):
output_ins
=
outputs
[
indices
[:,
0
]
+
b
,
indices
[:,
1
]
+
j
,
indices
[:,
2
],
indices
[:,
3
]]
mean_feat
=
torch
.
mean
(
output_ins
).
expand_as
(
output_ins
)
outputs_mean
[
indices
[:,
0
]
+
b
,
indices
[:,
1
]
+
j
,
indices
[:,
2
],
indices
[:,
3
]]
=
mean_feat
return
outputs_mean
class
MultiscaleDiscriminator
(
nn
.
Module
):
def
__init__
(
self
,
input_nc
,
ndf
=
64
,
n_layers
=
3
,
norm_layer
=
nn
.
BatchNorm2d
,
use_sigmoid
=
False
,
num_D
=
3
,
getIntermFeat
=
False
):
super
(
MultiscaleDiscriminator
,
self
).
__init__
()
self
.
num_D
=
num_D
self
.
n_layers
=
n_layers
self
.
getIntermFeat
=
getIntermFeat
for
i
in
range
(
num_D
):
netD
=
NLayerDiscriminator
(
input_nc
,
ndf
,
n_layers
,
norm_layer
,
use_sigmoid
,
getIntermFeat
)
if
getIntermFeat
:
for
j
in
range
(
n_layers
+
2
):
setattr
(
self
,
'scale'
+
str
(
i
)
+
'_layer'
+
str
(
j
),
getattr
(
netD
,
'model'
+
str
(
j
)))
else
:
setattr
(
self
,
'layer'
+
str
(
i
),
netD
.
model
)
self
.
downsample
=
nn
.
AvgPool2d
(
3
,
stride
=
2
,
padding
=
[
1
,
1
],
count_include_pad
=
False
)
def
singleD_forward
(
self
,
model
,
input
):
if
self
.
getIntermFeat
:
result
=
[
input
]
for
i
in
range
(
len
(
model
)):
result
.
append
(
model
[
i
](
result
[
-
1
]))
return
result
[
1
:]
else
:
return
[
model
(
input
)]
def
forward
(
self
,
input
):
num_D
=
self
.
num_D
result
=
[]
input_downsampled
=
input
for
i
in
range
(
num_D
):
if
self
.
getIntermFeat
:
model
=
[
getattr
(
self
,
'scale'
+
str
(
num_D
-
1
-
i
)
+
'_layer'
+
str
(
j
))
for
j
in
range
(
self
.
n_layers
+
2
)]
else
:
model
=
getattr
(
self
,
'layer'
+
str
(
num_D
-
1
-
i
))
result
.
append
(
self
.
singleD_forward
(
model
,
input_downsampled
))
if
i
!=
(
num_D
-
1
):
input_downsampled
=
self
.
downsample
(
input_downsampled
)
return
result
# Defines the PatchGAN discriminator with the specified arguments.
class
NLayerDiscriminator
(
nn
.
Module
):
def
__init__
(
self
,
input_nc
,
ndf
=
64
,
n_layers
=
3
,
norm_layer
=
nn
.
BatchNorm2d
,
use_sigmoid
=
False
,
getIntermFeat
=
False
):
super
(
NLayerDiscriminator
,
self
).
__init__
()
self
.
getIntermFeat
=
getIntermFeat
self
.
n_layers
=
n_layers
kw
=
4
padw
=
int
(
np
.
ceil
((
kw
-
1.0
)
/
2
))
sequence
=
[[
nn
.
Conv2d
(
input_nc
,
ndf
,
kernel_size
=
kw
,
stride
=
2
,
padding
=
padw
),
nn
.
LeakyReLU
(
0.2
,
True
)]]
nf
=
ndf
for
n
in
range
(
1
,
n_layers
):
nf_prev
=
nf
nf
=
min
(
nf
*
2
,
512
)
sequence
+=
[[
nn
.
Conv2d
(
nf_prev
,
nf
,
kernel_size
=
kw
,
stride
=
2
,
padding
=
padw
),
norm_layer
(
nf
),
nn
.
LeakyReLU
(
0.2
,
True
)
]]
nf_prev
=
nf
nf
=
min
(
nf
*
2
,
512
)
sequence
+=
[[
nn
.
Conv2d
(
nf_prev
,
nf
,
kernel_size
=
kw
,
stride
=
1
,
padding
=
padw
),
norm_layer
(
nf
),
nn
.
LeakyReLU
(
0.2
,
True
)
]]
sequence
+=
[[
nn
.
Conv2d
(
nf
,
1
,
kernel_size
=
kw
,
stride
=
1
,
padding
=
padw
)]]
if
use_sigmoid
:
sequence
+=
[[
nn
.
Sigmoid
()]]
if
getIntermFeat
:
for
n
in
range
(
len
(
sequence
)):
setattr
(
self
,
'model'
+
str
(
n
),
nn
.
Sequential
(
*
sequence
[
n
]))
else
:
sequence_stream
=
[]
for
n
in
range
(
len
(
sequence
)):
sequence_stream
+=
sequence
[
n
]
self
.
model
=
nn
.
Sequential
(
*
sequence_stream
)
def
forward
(
self
,
input
):
if
self
.
getIntermFeat
:
res
=
[
input
]
for
n
in
range
(
self
.
n_layers
+
2
):
model
=
getattr
(
self
,
'model'
+
str
(
n
))
res
.
append
(
model
(
res
[
-
1
]))
return
res
[
1
:]
else
:
return
self
.
model
(
input
)
from
torchvision
import
models
class
Vgg19
(
torch
.
nn
.
Module
):
def
__init__
(
self
,
requires_grad
=
False
):
super
(
Vgg19
,
self
).
__init__
()
vgg_pretrained_features
=
models
.
vgg19
(
pretrained
=
True
).
features
self
.
slice1
=
torch
.
nn
.
Sequential
()
self
.
slice2
=
torch
.
nn
.
Sequential
()
self
.
slice3
=
torch
.
nn
.
Sequential
()
self
.
slice4
=
torch
.
nn
.
Sequential
()
self
.
slice5
=
torch
.
nn
.
Sequential
()
for
x
in
range
(
2
):
self
.
slice1
.
add_module
(
str
(
x
),
vgg_pretrained_features
[
x
])
for
x
in
range
(
2
,
7
):
self
.
slice2
.
add_module
(
str
(
x
),
vgg_pretrained_features
[
x
])
for
x
in
range
(
7
,
12
):
self
.
slice3
.
add_module
(
str
(
x
),
vgg_pretrained_features
[
x
])
for
x
in
range
(
12
,
21
):
self
.
slice4
.
add_module
(
str
(
x
),
vgg_pretrained_features
[
x
])
for
x
in
range
(
21
,
30
):
self
.
slice5
.
add_module
(
str
(
x
),
vgg_pretrained_features
[
x
])
if
not
requires_grad
:
for
param
in
self
.
parameters
():
param
.
requires_grad
=
False
def
forward
(
self
,
X
):
h_relu1
=
self
.
slice1
(
X
)
h_relu2
=
self
.
slice2
(
h_relu1
)
h_relu3
=
self
.
slice3
(
h_relu2
)
h_relu4
=
self
.
slice4
(
h_relu3
)
h_relu5
=
self
.
slice5
(
h_relu4
)
out
=
[
h_relu1
,
h_relu2
,
h_relu3
,
h_relu4
,
h_relu5
]
return
out
models/runmodel.py
浏览文件 @
cde29be4
...
...
@@ -5,7 +5,6 @@ from util import mosaic
from
util
import
data
import
torch
def
run_unet
(
img
,
net
,
size
=
128
,
use_gpu
=
True
):
img
=
impro
.
image2folat
(
img
,
3
)
img
=
img
.
reshape
(
1
,
3
,
size
,
size
)
...
...
@@ -25,9 +24,12 @@ def run_unet_rectim(img,net,size = 128,use_gpu = True):
mask
=
impro
.
mergeimage
(
mask1
,
mask2
,
img
)
return
mask
def
run_pix2pix
(
img
,
net
,
size
=
128
,
use_gpu
=
True
):
img
=
impro
.
resize
(
img
,
size
)
img
=
data
.
im2tensor
(
img
,
use_gpu
=
use_gpu
)
def
run_pix2pix
(
img
,
net
,
opt
):
if
opt
.
netG
==
'HD'
:
img
=
impro
.
resize
(
img
,
512
)
else
:
img
=
impro
.
resize
(
img
,
128
)
img
=
data
.
im2tensor
(
img
,
use_gpu
=
opt
.
use_gpu
)
img_fake
=
net
(
img
)
img_fake
=
data
.
tensor2im
(
img_fake
)
return
img_fake
...
...
@@ -38,7 +40,6 @@ def get_ROI_position(img,net,opt):
x
,
y
,
halfsize
,
area
=
impro
.
boundingSquare
(
mask
,
1
)
return
mask
,
x
,
y
,
area
def
get_mosaic_position
(
img_origin
,
net_mosaic_pos
,
opt
):
mask
=
run_unet_rectim
(
img_origin
,
net_mosaic_pos
,
use_gpu
=
opt
.
use_gpu
)
mask
=
impro
.
mask_threshold
(
mask
,
10
,
128
)
...
...
models/unet_parts.py
浏览文件 @
cde29be4
...
...
@@ -54,7 +54,7 @@ class Upsample(nn.Module):
def
forward
(
self
,
x
):
return
F
.
interpolate
(
x
,
scale_factor
=
self
.
scale_factor
,
mode
=
'bilinear'
,
align_corners
=
True
)
F
.
interpolate
class
up
(
nn
.
Module
):
def
__init__
(
self
,
in_ch
,
out_ch
,
bilinear
=
True
):
super
(
up
,
self
).
__init__
()
...
...
options.py
浏览文件 @
cde29be4
...
...
@@ -23,8 +23,8 @@ class Options():
self
.
parser
.
add_argument
(
'--mask_threshold'
,
type
=
int
,
default
=
64
,
help
=
'threshold of recognize mosaic position 0~255'
)
self
.
parser
.
add_argument
(
'--output_size'
,
type
=
int
,
default
=
0
,
help
=
'size of output file,if 0 -> origin'
)
#
Add
Mosaic
self
.
parser
.
add_argument
(
'--netG'
,
type
=
str
,
default
=
'auto'
,
help
=
'select model to use for netG(clean mosaic) -> auto | unet_128 | resnet_9blocks'
)
#
Clean
Mosaic
self
.
parser
.
add_argument
(
'--netG'
,
type
=
str
,
default
=
'auto'
,
help
=
'select model to use for netG(clean mosaic) -> auto | unet_128 | resnet_9blocks
| HD
'
)
self
.
parser
.
add_argument
(
'--mosaic_position_model_path'
,
type
=
str
,
default
=
'auto'
,
help
=
'name of model use to find mosaic position'
)
self
.
parser
.
add_argument
(
'--no_feather'
,
action
=
'store_true'
,
help
=
'if true, no edge feather,but run faster'
)
self
.
parser
.
add_argument
(
'--medfilt_num'
,
type
=
int
,
default
=
11
,
help
=
'medfilt window of mosaic movement in the video'
)
...
...
@@ -36,11 +36,16 @@ class Options():
self
.
initialize
()
self
.
opt
=
self
.
parser
.
parse_args
()
if
self
.
opt
.
netG
==
'auto'
:
if
self
.
opt
.
netG
==
'auto'
and
self
.
opt
.
mode
==
'clean'
:
if
'unet_128'
in
self
.
opt
.
model_path
:
self
.
opt
.
netG
=
'unet_128'
elif
'resnet_9blocks'
in
self
.
opt
.
model_path
:
self
.
opt
.
netG
=
'resnet_9blocks'
elif
'HD'
in
self
.
opt
.
model_path
:
self
.
opt
.
netG
=
'HD'
else
:
print
(
'Type of Generator error!'
)
if
self
.
opt
.
mosaic_position_model_path
==
'auto'
:
_path
=
os
.
path
.
join
(
os
.
path
.
split
(
self
.
opt
.
model_path
)[
0
],
'mosaic_position.pth'
)
...
...
util/data.py
浏览文件 @
cde29be4
...
...
@@ -20,11 +20,12 @@ def tensor2im(image_tensor, imtype=np.uint8, rgb2bgr = True):
def
im2tensor
(
image_numpy
,
imtype
=
np
.
uint8
,
bgr2rgb
=
True
,
reshape
=
True
,
use_gpu
=
True
):
h
,
w
=
image_numpy
.
shape
[:
2
]
if
bgr2rgb
:
image_numpy
=
image_numpy
[...,::
-
1
]
-
np
.
zeros_like
(
image_numpy
)
image_tensor
=
transform
(
image_numpy
)
if
reshape
:
image_tensor
=
image_tensor
.
reshape
(
1
,
3
,
128
,
128
)
image_tensor
=
image_tensor
.
reshape
(
1
,
3
,
h
,
w
)
if
use_gpu
:
image_tensor
=
image_tensor
.
cuda
()
return
image_tensor
\ No newline at end of file
util/mosaic.py
浏览文件 @
cde29be4
...
...
@@ -55,12 +55,16 @@ def addmosaic_normal(img,mask,n,out_size = 0,model = 'squa_avg'):
return
img_mosaic
def
addmosaic_autosize
(
img
,
mask
,
model
):
def
addmosaic_autosize
(
img
,
mask
,
model
,
area_type
=
'normal'
):
h
,
w
=
img
.
shape
[:
2
]
mask
=
cv2
.
resize
(
mask
,(
w
,
h
))
alpha
=
np
.
min
((
w
,
h
))
/
512
try
:
area
=
mask_area
(
mask
)
if
area_type
==
'normal'
:
area
=
mask_area
(
mask
)
elif
area_type
==
'bounding'
:
w
,
h
=
cv2
.
boundingRect
(
mask
)[
2
:]
area
=
w
*
h
except
:
area
=
0
area
=
area
/
(
alpha
*
alpha
)
...
...
@@ -76,19 +80,23 @@ def addmosaic_autosize(img,mask,model):
pass
return
img_mosaic
def
addmosaic_random
(
img
,
mask
):
def
addmosaic_random
(
img
,
mask
,
area_type
=
'normal'
):
# img = resize(img,512)
h
,
w
=
img
.
shape
[:
2
]
mask
=
cv2
.
resize
(
mask
,(
w
,
h
))
alpha
=
np
.
min
((
w
,
h
))
/
512
#area_avg=5925*4
try
:
area
=
mask_area
(
mask
)
if
area_type
==
'normal'
:
area
=
mask_area
(
mask
)
elif
area_type
==
'bounding'
:
w
,
h
=
cv2
.
boundingRect
(
mask
)[
2
:]
area
=
w
*
h
except
:
area
=
0
area
=
area
/
(
alpha
*
alpha
)
if
area
>
50000
:
img_mosaic
=
random_mod
(
img
,
mask
,
alpha
*
random
.
uniform
(
16
,
28
))
img_mosaic
=
random_mod
(
img
,
mask
,
alpha
*
random
.
uniform
(
16
,
30
))
elif
20000
<
area
<=
50000
:
img_mosaic
=
random_mod
(
img
,
mask
,
alpha
*
random
.
uniform
(
12
,
20
))
elif
5000
<
area
<=
20000
:
...
...
util/util.py
浏览文件 @
cde29be4
import
os
import
shutil
def
Traversal
(
filedir
):
file_list
=
[]
for
root
,
dirs
,
files
in
os
.
walk
(
filedir
):
...
...
@@ -29,6 +30,13 @@ def writelog(path,log):
f
=
open
(
path
,
'a+'
)
f
.
write
(
log
+
'
\n
'
)
def
makedirs
(
path
):
if
os
.
path
.
isdir
(
path
):
print
(
path
,
'existed'
)
else
:
os
.
makedirs
(
path
)
print
(
'makedir:'
,
path
)
def
clean_tempfiles
(
tmp_init
=
True
):
if
os
.
path
.
isdir
(
'./tmp'
):
shutil
.
rmtree
(
'./tmp'
)
...
...
编辑
预览
Markdown
is supported
0%
请重试
或
添加新附件
.
添加附件
取消
You are about to add
0
people
to the discussion. Proceed with caution.
先完成此消息的编辑!
取消
想要评论请
注册
或
登录