def train(self):
# Prepare image pyramid
train_img = read_img(self.config)
real = resize_img(train_img, self.config.start_scale, self.config)
self.reals = creat_reals_pyramid(real, self.reals, self.config)
prev_nfc = 0
self.writer = SummaryWriter(f'{self.config.exp_dir}/logs')
# Pyramid training
for scale_iter in range(self.config.stop_scale + 1):
# Become larger as scale_iter increase (maximum=128)
self.config.nfc = min(
self.config.nfc_init * pow(2, math.floor(scale_iter / 4)), 128)
self.config.min_nfc = min(
self.config.min_nfc_init * pow(2, math.floor(scale_iter / 4)),
128)
# Prepare directory to save images
self.config.result_dir = f'{self.config.exp_dir}/{scale_iter}'
os.makedirs(self.config.result_dir, exist_ok=True)
plt.imsave(f'{self.config.result_dir}/real_scale.png',
torch2np(self.reals[scale_iter]),
vmin=0,
vmax=1)
cur_discriminator, cur_generator = self.init_models()
if prev_nfc == self.config.nfc:
cur_generator.load_state_dict(
torch.load(
f'{self.config.exp_dir}/{scale_iter - 1}/generator.pth'
))
cur_discriminator.load_state_dict(
torch.load(
f'{self.config.exp_dir}/{scale_iter - 1}/discriminator.pth'
))
cur_z, cur_generator = self.train_single_stage(
cur_discriminator, cur_generator)
cur_generator = reset_grads(cur_generator, False)
cur_generator.eval()
cur_discriminator = reset_grads(cur_discriminator, False)
cur_discriminator.eval()
self.Gs.append(cur_generator)
self.Zs.append(cur_z)
self.noise_amps.append(self.config.noise_amp)
torch.save(self.Zs, f'{self.config.exp_dir}/Zs.pth')
torch.save(self.Gs, f'{self.config.exp_dir}/Gs.pth')
torch.save(self.reals, f'{self.config.exp_dir}/reals.pth')
torch.save(self.noise_amps, f'{self.config.exp_dir}/noiseAmp.pth')
prev_nfc = self.config.nfc
del cur_discriminator, cur_generator
return
def anno_inspect():
with open('/media/csh/data/coco_2017/annotations/instances_val2017.json'
) as f:
res = json.load(f)
res = res['annotations']
coco_anno = "/media/csh/data/coco_2017/val2017"
for anno in res[:10]:
image_id = anno['image_id']
img_path = os.path.join(coco_anno, "%012d.jpg" % image_id)
bbox = np.array(list(map(lambda x: int(x), anno['bbox'])))
bbox[..., 2] = bbox[..., 0] + bbox[..., 2]
bbox[..., 3] = bbox[..., 1] + bbox[..., 3]
img = read_img(img_path)
print(image_id, bbox)
draw_box(img, bbox, debug=True)
def interactive_inspect(self, show_anno=False, imgs=None):
if imgs is None:
img_path = self.task_util.get_eval_imgs()
else:
img_path = imgs
for p in img_path:
print(p)
image = read_img(p)
image_id = os.path.splitext(os.path.basename(p))[0].lstrip('0')
anno_box, anno_human_label = self.task_util.get_image_anno(
image_id)
predict_box, predict_box_confidence, predict_label, predict_score = self.predict(
image)
predict_human_label = [
self.task_util.label_classes[x] for x in predict_label
]
show_tags = list(
zip(predict_human_label, [
"%.2f %.2f" % (box_conf, label_score) for box_conf,
label_score in zip(predict_box_confidence, predict_score)
]))
print('predict labels: %s' % str(show_tags))
print('annotate label: %s' % anno_human_label)
if show_anno and len(anno_box) != 0:
draw_box(image.copy(),
np.array(anno_box),
debug=True,
title='anno',
text=anno_human_label)
if len(predict_box) != 0:
draw_box(image.copy(),
predict_box,
debug=True,
title='predict',
text=show_tags)
def predict(self, image, specific=True, sess: tf.compat.v1.Session = None):
if sess is None:
sess = get_session()
if isinstance(image, str):
origin_image = read_img(image)
else:
origin_image = image
origin_h, origin_w = origin_image.shape[:-1]
# resize image and return the image with model input size and offset and scale infomation
input_image, offset_xy, scale_xy = _image_preprocess(
origin_image, input_shape=self.input_shape)
input_image = np.expand_dims(input_image, 0)
# box [[xmin, ymin, xmax, ymax]...]
# box scores [s1, s2, s3...] 0 <= s_i <= 1
# labels etc [c1, c2, c3 ...] 0 <= c_i <= C
# score etc [t1, t2, t3 ...] 0 <= t_i <= 1
b, bc, l, ls = sess.run(
[self._box, self._box_confidence, self._labels, self._scores],
feed_dict={self.input_image: input_image})
b[:, [0, 2]] = b[:, [0, 2]] - offset_xy[0]
b[:, [1, 3]] = b[:, [1, 3]] - offset_xy[1]
b[:, [0, 2]] = b[:, [0, 2]] / scale_xy[0]
b[:, [1, 3]] = b[:, [1, 3]] / scale_xy[1]
b[:, [0, 2]] = np.clip(b[:, [0, 2]], 0, origin_w - 1)
b[:, [1, 3]] = np.clip(b[:, [1, 3]], 0, origin_h - 1)
# nms
_box, _box_confidence, _label, _score = _bbox_proprocess_soft(
b, bc, l, ls, image=origin_image)
if specific:
return _box, _box_confidence, _label, _score
else:
return _box, _label, _score
import matplotlib.pyplot as plt
from config import Config
# from model.ACM_SinGAN import SinGAN_ACM as SinGAN
from model.SinGAN import SinGAN
from utils.image import read_img, torch2np
from utils.utils import process_config, adjust_scales, calcul_sr_scale
if __name__ == '__main__':
process_config(Config)
Config.infer_dir = f'{Config.exp_dir}/infer_pyramid' if Config.save_all_pyramid else f'{Config.exp_dir}/infer'
singan = SinGAN(config=Config)
singan.load_trained_weights()
inference_img = read_img(Config)
os.makedirs(Config.infer_dir, exist_ok=True)
plt.imsave(f'{Config.exp_dir}/real.png',
torch2np(inference_img),
vmin=0,
vmax=1)
if Config.mode == 'train':
adjust_scales(inference_img, Config)
start_img_input = singan.create_inference_input()
elif Config.mode == 'train_SR':
in_scale, iter_num = calcul_sr_scale(Config)
Config.scale_factor = 1 / in_scale
Config.scale_factor_init = 1 / in_scale
Config.scale_w = 1
Config.scale_h = 1
import os
from shutil import copyfile
from config import Config
from model.SinGAN import SinGAN
from model.ACM_SinGAN import SinGAN_ACM
from utils.image import read_img
from utils.utils import process_config, adjust_scales, calcul_sr_scale
# All data : [B C H W]
if __name__ == '__main__':
os.makedirs(Config.exp_dir, exist_ok=True)
copyfile('config.py', f'{Config.exp_dir}/config.py')
process_config(Config)
train_img = read_img(Config)
if Config.mode == "train":
adjust_scales(train_img, Config)
elif Config.mode == "train_SR":
calcul_sr_scale(Config)
adjust_scales(train_img, Config)
singan = SinGAN_ACM(config=Config) if Config.use_acm else SinGAN(
config=Config)
singan.train()
bbox = np.array(list(map(lambda x: int(x), anno['bbox'])))
bbox[..., 2] = bbox[..., 0] + bbox[..., 2]
bbox[..., 3] = bbox[..., 1] + bbox[..., 3]
img = read_img(img_path)
print(image_id, bbox)
draw_box(img, bbox, debug=True)
if __name__ == '__main__':
image_dir = "/media/csh/data/coco_2017/val2017"
anno_dir = "/media/csh/data/coco_2017/annotations/instances_val2017.json"
with open(anno_dir) as f:
coco_anno = json.load(f)
for anno in coco_anno['annotations']:
image_path = os.path.join(image_dir,
"%012d" % anno['image_id']) + '.jpg'
bbox = anno['bbox']
break
img = read_img(image_path)
bbox = xywh2xymimx(bbox)
draw_box(img, bbox, debug=True)
# test_img = ["/media/csh/data/coco_2017/val2017/000000477118.jpg",
# "/media/csh/data/coco_2017/val2017/000000006818.jpg",
# "/media/csh/data/coco_2017/val2017/000000004765.jpg"]
# for t_img in test_img:
# img = read_img(t_img)
# letterbox_resize(img, (416, 416), show=True)