class FaceSegementer():
def __init__(self):
self.model = BiSeNet(n_classes=19)
self.model.cuda()
weights_filename = 'face-parsing-02dd3f6f.pth'
url = 'https://rewriting.csail.mit.edu/data/models/' + weights_filename
try:
sd = torch.hub.load_state_dict_from_url(url) # pytorch 1.1+
except Exception:
sd = torch.hub.model_zoo.load_url(url) # pytorch 1.0
self.model.load_state_dict(sd)
self.model.eval()
def segment_batch(self, xs):
results = []
og_size = xs.shape[2:]
xs = F.interpolate(xs, size=(512, 512))
for x in xs:
x = x.unsqueeze(dim=0)
out = self.model(x)[0]
mask = torch.from_numpy(
out.squeeze(0).cpu().numpy().argmax(0)).unsqueeze(dim=0)
results.append(mask)
masks = torch.stack(results).float()
masks = F.interpolate(masks, size=og_size).long()
return masks
def evaluate(dspth='./', cp=''):
n_classes = 19
net = BiSeNet(n_classes=n_classes)
net.cuda()
save_pth = cp
print('model : {}'.format(save_pth))
net.load_state_dict(torch.load(save_pth))
net.eval()
to_tensor = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),
])
with torch.no_grad():
idx = 0
for image_path in os.listdir(dspth):
img = Image.open(osp.join(dspth, image_path))
image = img.resize((512, 512), Image.BILINEAR)
img = to_tensor(image)
img = torch.unsqueeze(img, 0)
img = img.cuda()
out = net(img)[0]
parsing = out.squeeze(0).cpu().numpy().argmax(0)
# print(parsing)
idx += 1
print('<{}> image : '.format(idx), np.unique(parsing))
test_result = './test_result'
if not osp.exists(test_result):
os.makedirs(test_result)
vis_parsing_maps(image,
parsing,
stride=1,
save_im=True,
save_path=osp.join(test_result, image_path))
def evaluate(cfg,
respth='./res/test_res',
dspth='./data',
cp='model_final_diss.pth'):
if not os.path.exists(respth):
os.makedirs(respth)
n_classes = 19
net = BiSeNet(n_classes=n_classes)
if cfg["GPU"]["enable"]:
net.cuda(cfg["GPU"]["name"])
net.load_state_dict(torch.load(cp))
net.eval()
to_tensor = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),
])
with torch.no_grad():
for image_path in os.listdir(dspth):
img = Image.open(osp.join(dspth, image_path))
image = img.resize((512, 512), Image.BILINEAR)
img = to_tensor(image)
img = torch.unsqueeze(img, 0)
img = img.cuda(cfg["GPU"]["name"]) if cfg["GPU"]["enable"] else img
out = net(img)[0]
parsing = out.squeeze(0).cpu().numpy().argmax(0)
parsing = \
torch.nn.functional.interpolate(torch.tensor(np.expand_dims(parsing, axis=(0, 1))).type(torch.tensor),
size=(cfg["mask_size"], cfg["mask_size"]))[0][0].numpy()
hkl.dump(parsing, join(respth, image_path[:-4]))
def evaluate(image_path='./imgs/116.jpg', cp='cp/79999_iter.pth'):
# if not os.path.exists(respth):
# os.makedirs(respth)
n_classes = 19
net = BiSeNet(n_classes=n_classes) # BiSeNet으로 net이리는 인스턴스 생성됨. 인자로 19 넣어서 만듦.
net.cuda() # Tensor들을 GPU로 보내기
net.load_state_dict(torch.load(cp))
net.eval()
to_tensor = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),
])
with torch.no_grad():
img = Image.open(image_path)
image = img.resize((512, 512), Image.BILINEAR)
img = to_tensor(image)
img = torch.unsqueeze(img, 0)
img = img.cuda()
out = net(img)[0]
parsing = out.squeeze(0).cpu().numpy().argmax(0)
# print(np.shape(img),np.shape(parsing),parsing) => torch.Size([1,3,512, 512]) / (512, 512) / [0 0 0...0 0 0]~[17 17 17... 17 17 17]
# print(np.unique(parsing))
# vis_parsing_maps(image, parsing, stride=1, save_im=False, save_path=osp.join(respth, dspth))
return parsing
def evaluate(image_path='./images/img1.jpg', cp='79999_iter.pth'):
# if not os.path.exists(respth):
# os.makedirs(respth)
n_classes = 19
net = BiSeNet(n_classes=n_classes)
#net.cuda() ## uncomment this
#net.load_state_dict(torch.load(cp) with map_location=torch.device('cpu')) ##net.load_state_dict(torch.load(cp))
net.load_state_dict(torch.load(cp, map_location='cpu'))
net.eval()
to_tensor = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),
])
with torch.no_grad():
img = Image.open(image_path)
image = img.resize((512, 512), Image.BILINEAR)
img = to_tensor(image)
img = torch.unsqueeze(img, 0)
#img = img.cuda() ## uncomment this
out = net(img)[0]
parsing = out.squeeze(0).cpu().numpy().argmax(0)
# print(parsing)
# print(np.unique(parsing))
# vis_parsing_maps(image, parsing, stride=1, save_im=False, save_path=osp.join(respth, dspth))
return parsing
def evaluate(respth='./res', dspth='./data'):
## logger
logger = logging.getLogger()
## model
logger.info('\n')
logger.info('====' * 20)
logger.info('evaluating the model ...\n')
logger.info('setup and restore model')
n_classes = 19
net = BiSeNet(n_classes=n_classes)
save_pth = osp.join(respth, 'model_final.pth')
net.load_state_dict(torch.load(save_pth))
net.cuda()
net.eval()
## dataset
batchsize = 5
n_workers = 2
dsval = CityScapes(dspth, mode='val')
dl = DataLoader(dsval,
batch_size=batchsize,
shuffle=False,
num_workers=n_workers,
drop_last=False)
## evaluator
logger.info('compute the mIOU')
evaluator = MscEval(net, dl)
## eval
mIOU = evaluator.evaluate()
logger.info('mIOU is: {:.6f}'.format(mIOU))
def evaluate(image_path="./imgs/116.jpg", cp="cp/79999_iter.pth"):
# if not os.path.exists(respth):
# os.makedirs(respth)
n_classes = 19
net = BiSeNet(n_classes=n_classes)
if CUDA:
net.cuda()
net.load_state_dict(torch.load(cp, map_location=torch.device("cpu")))
net.eval()
to_tensor = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),
])
with torch.no_grad():
#img = Image.fromarray(image_path)
img = Image.open(image_path)
image = img.resize((512, 512), Image.BILINEAR)
img = to_tensor(image)
img = torch.unsqueeze(img, 0)
if CUDA:
img = img.cuda()
out = net(img)[0]
parsing = out.squeeze(0).cpu().numpy().argmax(0)
print(parsing)
print(np.unique(parsing))
vis_parsing_maps(image, parsing, stride=1)
return parsing
def evaluate(respth='./res/test_res',
dspth='./data',
cp='model_final_diss.pth'):
if not os.path.exists(respth):
os.makedirs(respth)
n_classes = 19
net = BiSeNet(n_classes=n_classes)
net.cuda()
save_pth = osp.join('res/cp', cp)
net.load_state_dict(torch.load(save_pth))
net.eval()
to_tensor = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),
])
with torch.no_grad():
for image_path in os.listdir(dspth):
img = Image.open(osp.join(dspth, image_path))
image = img.resize((512, 512), Image.BILINEAR)
img = to_tensor(image)
img = torch.unsqueeze(img, 0)
img = img.cuda()
out = net(img)[0]
parsing = out.squeeze(0).cpu().numpy().argmax(0)
# print(parsing)
print(np.unique(parsing))
vis_parsing_maps(image,
parsing,
stride=1,
save_im=True,
save_path=osp.join(respth, image_path))
def __init__(self):
sys.path.append('resources/face_parsing_pytorch/')
from model import BiSeNet
n_classes = 19
seg_net = BiSeNet(n_classes=n_classes).cuda()
save_pth = 'resources/face_parsing_pytorch/res/cp/79999_iter.pth'
seg_net.load_state_dict(torch.load(save_pth))
seg_net.eval()
self.seg_net = seg_net
def evaluate(respth='./results/data_src',
dspth='../data',
cp='79999_iter.pth'):
respth = osp.join(os.path.abspath(os.path.dirname(__file__)), respth)
if not os.path.exists(respth):
os.makedirs(respth)
n_classes = 19
net = BiSeNet(n_classes=n_classes)
net.cuda()
model_path = osp.join(os.path.abspath(os.path.dirname(__file__)), cp)
data_path = osp.join(os.path.abspath(os.path.dirname(__file__)), dspth)
net.load_state_dict(torch.load(model_path))
net.eval()
to_tensor = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),
])
face_detector = FaceDetector(device="cuda", verbose=False)
with torch.no_grad():
for image_path in os.listdir(data_path):
img = Image.open(osp.join(data_path, image_path))
image = img.resize((512, 512), Image.BILINEAR)
# --------------------------------------------------------------------------------------
bgr_image = cv2.imread(osp.join(data_path, image_path))
detected_faces = face_detector.detect_from_image(
cv2.cvtColor(bgr_image, cv2.COLOR_BGR2RGB))
lx, ly, rx, ry = int(detected_faces[0][0]), int(
detected_faces[0][1]), int(detected_faces[0][2]), int(
detected_faces[0][3])
face = bgr_image[ly:ry, lx:rx, :].copy()
# upscale the image
face = cv2.resize(face, (0, 0), fx=5, fy=5)
# cv2.imwrite(osp.join(respth, image_path), face)
face_rgb = cv2.cvtColor(face, cv2.COLOR_BGR2RGB)
# --------------------------------------------------------------------------------------
img = to_tensor(face_rgb)
img = torch.unsqueeze(img, 0)
img = img.cuda()
out = net(img)[0]
parsing = out.squeeze(0).cpu().numpy().argmax(0)
print(np.unique(parsing))
out_face_mask = vis_parsing_maps(face_rgb,
parsing,
stride=1,
save_im=True,
save_path=osp.join(
respth, image_path))
origin_face_mask = cv2.resize(out_face_mask, (rx - lx, ry - ly))
image_mask = np.zeros(np.shape(bgr_image)[:2], dtype=np.uint8)
image_mask[ly:ry, lx:rx] = origin_face_mask
result = merge(bgr_image, image_mask)
cv2.imwrite(osp.join(respth, image_path)[:-4] + '.png', result)
def setup_face_parsing():
global parser
parser_file = downloads.download_from_gdrive(
gdrive_fileid='154JgKpzCPW82qINcVieuPH3fZ2e0P812',
output_path='face-parsing/79999_iter.pth')
n_classes = 19
parser = BiSeNet(n_classes=n_classes)
parser.cuda()
parser.load_state_dict(torch.load(parser_file))
parser.eval()
def create_faceparse_model(path_model):
n_classes = 19
net = BiSeNet(n_classes=n_classes)
net.cuda()
print('model : {}'.format(path_model))
net.load_state_dict(torch.load(path_model))
net.eval()
return net
def buildnet(cp='79999_iter.pth'):
n_classes = 19
net = BiSeNet(n_classes=n_classes)
# net.cuda()
save_pth = osp.join('res/cp', cp)
net.load_state_dict(torch.load(save_pth, map_location="cpu"))
net.eval()
to_tensor = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),
])
return net, to_tensor
def main(haar_cascade_filepath, state_dict):
app = QtWidgets.QApplication(sys.argv)
n_classes = 19
net = BiSeNet(
n_classes=n_classes) # BiSeNet으로 net이리는 인스턴스 생성됨. 인자로 19 넣어서 만듦.
net.cuda() # Tensor들을 GPU로 보내기
net.load_state_dict(state_dict)
net.eval()
main_window = MainWindow()
main_widget = MainWidget(haar_cascade_filepath, net, main_window)
main_window.setCentralWidget(main_widget)
main_window.show()
sys.exit(app.exec_())
def parser(image_path='./imgs/116.jpg', cp='cp/79999_iter.pth'):
n_classes = 19
net = BiSeNet(n_classes=n_classes)
device = torch.device("cpu")
net.to(device)
net.load_state_dict(torch.load(cp,map_location='cpu'))
net.eval()
to_tensor = transforms.Compose([transforms.ToTensor(),
transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),])
with torch.no_grad():
# img = Image.open(image_path)
img = to_tensor(im)
img = torch.unsqueeze(img, 0)
out = net(img)[0]
parsing = out.squeeze(0).numpy().argmax(0)
return parsing
def evaluate(face_img, cp='79999_iter.pth'):
n_classes = 19
net = BiSeNet(n_classes=n_classes)
# net.cuda()
save_pth = osp.join('network/', cp)
net.load_state_dict(torch.load(save_pth, map_location=torch.device('cpu')))
net.eval()
to_tensor = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),
])
with torch.no_grad():
img = to_tensor(face_img)
img = torch.unsqueeze(img, 0)
# img = img.cuda()
out = net(img)[0]
# print(out.shape) # torch.Size([1, 19, 512, 512])
parsing = out.squeeze(0).cpu().numpy().argmax(0)
return parsing.astype(np.uint8)
def label_images(images: torch.tensor, cp='model_final_diss.pth'):
"""
Labels images
Args:
images: a (n x H x W x C) tensor with the images
Returns:
a (n x H x W) numpy array with the class labels
"""
n_classes = 19
net = BiSeNet(n_classes=n_classes)
net.cuda()
save_pth = os.path.join('res/cp', cp)
net.load_state_dict(torch.load(save_pth))
# need to move color channels dimension to the second dimension due to weird PIL transform issues
images = images.transpose(1, 3).transpose(2, 3)
transform_images = transforms.Compose([
transforms.ToPILImage(),
transforms.Resize((512, 512)),
transforms.ToTensor(),
transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),
])
net.eval()
with torch.no_grad():
labels = []
for image in tqdm.tqdm(images, desc="segmenting frames"):
img = transform_images(image)
img = torch.unsqueeze(img, 0)
img = img.cuda()
labels_probs = net(img)[0]
labels.append(labels_probs.squeeze(0).detach().cpu().argmax(0))
return torch.stack(labels, dim=0)
def getface(input_image):
save_pth = baseLoc+'face-parsing.PyTorch/79999_iter.pth'
input_image = Image.fromarray(input_image)
n_classes = 19
net = BiSeNet(n_classes=n_classes)
if torch.cuda.is_available():
net.cuda()
net.load_state_dict(torch.load(save_pth))
else:
net.load_state_dict(torch.load(save_pth, map_location=lambda storage, loc: storage))
net.eval()
to_tensor = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),
])
with torch.no_grad():
img = input_image.resize((512, 512), Image.BILINEAR)
img = to_tensor(img)
img = torch.unsqueeze(img, 0)
if torch.cuda.is_available():
img = img.cuda()
out = net(img)[0]
if torch.cuda.is_available():
parsing = out.squeeze(0).cpu().numpy().argmax(0)
else:
parsing = out.squeeze(0).numpy().argmax(0)
parsing = Image.fromarray(np.uint8(parsing))
parsing = parsing.resize(input_image.size)
parsing = np.array(parsing)
return parsing
def evaluate(image_path, cp):
n_classes = 19
net = BiSeNet(n_classes=n_classes)
net.cuda()
net.load_state_dict(torch.load(cp))
net.eval()
to_tensor = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),
])
with torch.no_grad():
img = Image.open(image_path)
image = img.resize((512, 512), Image.BILINEAR)
img = to_tensor(image)
img = torch.unsqueeze(img, 0)
img = img.cuda()
out = net(img)[0]
parsing = out.squeeze(0).cpu().numpy().argmax(0)
vis_parsing_maps(image, parsing, stride=1)
return parsing
def evaluate(respth='./res', dspth='./data'):
## logger
logger = logging.getLogger()
## model
logger.info('\n')
logger.info('====' * 20)
logger.info('evaluating the model ...\n')
logger.info('setup and restore model')
n_classes = 19
net = BiSeNet(n_classes=n_classes)
save_pth = osp.join(respth, 'model_final.pth')
net.load_state_dict(torch.load(save_pth))
net.cuda()
net.eval()
## dataset
batchsize = 2
n_workers = 1
dsval = CityScapes(dspth, mode='val')
dl = DataLoader(dsval,
batch_size=batchsize,
shuffle=False,
num_workers=n_workers,
drop_last=False)
## evaluator
logger.info('compute the mIOU')
evaluator = MscEval(net, dl)
# ## export to ONNX
# dummy_input = Variable(torch.randn(batchsize, 3, 1024, 1024)).cuda()
# onnx.export(net, dummy_input, "bisenet.proto", verbose=True)
## eval
mIOU = evaluator.evaluate()
logger.info('mIOU is: {:.6f}'.format(mIOU))
def exportImgAPI(
img,
style,
dspth='/home/KLTN_TheFaceOfArtFaceParsing/Updates/face_parsing/test-img',
cp='model_final_diss.pth'):
# img = Image.open(osp.join(dspth, image_path))
n_classes = 19
net = BiSeNet(n_classes=n_classes)
net.cuda()
save_pth = osp.join(
'/home/KLTN_TheFaceOfArtFaceParsing/Updates/face_parsing/res/cp', cp)
net.load_state_dict(torch.load(save_pth, map_location='cpu'))
net.eval()
to_tensor = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),
])
with torch.no_grad():
img = Image.fromarray(img)
image = img.resize((512, 512), Image.BILINEAR)
img = to_tensor(image)
img = torch.unsqueeze(img, 0)
img = img.cuda()
out = net(img)[0]
parsing = out.squeeze(0).cpu().numpy().argmax(0)
# print(parsing)
print(np.unique(parsing))
return vis_parsing_maps(image,
style,
parsing,
stride=1,
save_im=True,
save_path='')
import numpy as np
import torch
import math
from torch.utils.data import DataLoader
from model import BiSeNet
from face_dataset import FaceMask
from sklearn.metrics import f1_score
if __name__ == '__main__':
model_dir = '/home/jihyun/workspace/face_parsing/face-parsing.PyTorch/res/exp'
n_classes = 19
net = BiSeNet(n_classes=n_classes)
net.cuda()
net.eval()
data_root = '/home/jihyun/workspace/face_parsing/dataset/CelebAMask-HQ/'
cropsize = [448, 448]
n_img_per_gpu = 16
ds = FaceMask(data_root, cropsize=cropsize, mode='val')
dl = DataLoader(ds, batch_size=16, shuffle=False, drop_last=True)
f = open(os.path.join(model_dir, 'f_measure.txt'), 'w+')
for filename in os.listdir(os.path.join(model_dir, 'cp')):
if filename.endswith('70000_iter.pth'):
net.load_state_dict(
torch.load(os.path.join(model_dir, 'cp', filename)))
def testval(model_path, sv_dir='res', sv_pred=False):
n_classes = 17
confusion_matrix = np.zeros(
(n_classes, n_classes)) # num_classes x num_classes
cropsize = [448, 448]
data_root = '/home/data2/DATASET/vschallenge'
batch_size = 1
ds = FaceMask(data_root, cropsize=cropsize, mode='val')
dl = DataLoader(ds, batch_size=batch_size, shuffle=False, num_workers=0)
with torch.no_grad():
net = BiSeNet(n_classes=n_classes)
net.cuda()
net.load_state_dict(torch.load(model_path))
net.eval()
for iter, data in enumerate(dl):
im = data['img']
lb = data['label']
impth = data['impth']
lb = torch.squeeze(lb, 1)
im = im.cuda()
lb = lb.cuda()
out = net(im)[0]
size = lb.size()
pred = out.cpu().numpy().argmax(1)
gt = lb.cpu().numpy()
vis_parsing_maps(im.cpu(),
pred,
stride=1,
save_im=True,
save_path='res',
imspth=impth)
# vis_parsing_maps(im.cpu(), gt, stride=1, save_im=True, save_path='res_gt', imspth=impth)
confusion_matrix += get_confusion_matrix(
lb, out, size, n_classes, ignore=-1) # [16, 19, 448, 448]
if sv_pred:
sv_path = os.path.join(sv_dir, 'test_results')
if not os.path.exists(sv_path):
os.mkdir(sv_path)
# cv2.imwrite(sv_path+'/', ) add imname
if iter % 5 == 0:
pos = confusion_matrix.sum(1)
res = confusion_matrix.sum(0)
tp = np.diag(confusion_matrix)
pixel_acc = tp.sum() / pos.sum()
mean_acc = (tp / np.maximum(1.0, pos)).mean()
IoU_array = (tp / np.maximum(1.0, pos + res - tp))
mean_IoU = IoU_array.mean()
print(
'index/allimg {}/{}. mean_IoU: {:1.5f}. pixel_acc: {:1.5f}. mean_acc: {:1.5f}.'
.format(iter * batch_size,
len(dl) * batch_size, mean_IoU, pixel_acc,
mean_acc))
pos = confusion_matrix.sum(1)
res = confusion_matrix.sum(0)
tp = np.diag(confusion_matrix)
pixel_acc = tp.sum() / pos.sum()
mean_acc = (tp / np.maximum(1.0, pos)).mean()
IoU_array = (tp / np.maximum(1.0, pos + res - tp))
mean_IoU = IoU_array.mean()
print(
'mean_IoU: {:1.5f}. pixel_acc: {:1.5f}. mean_acc: {:1.5f}.'.format(
mean_IoU, pixel_acc, mean_acc))
return mean_IoU, IoU_array, pixel_acc, mean_acc
def train(opt):
# saving setting
opt.saved_path = opt.saved_path + opt.project
opt.log_path = os.path.join(opt.saved_path, 'tensorboard')
os.makedirs(opt.log_path, exist_ok=True)
os.makedirs(opt.saved_path, exist_ok=True)
# gpu setting
os.environ["CUDA_VISIBLE_DEVICES"] = '2, 3, 4, 5, 6'
gpu_number = torch.cuda.device_count()
# dataset setting
n_classes = 17
n_img_all_gpu = opt.batch_size * gpu_number
cropsize = [448, 448]
data_root = '/home/data2/DATASET/vschallenge'
num_workers = opt.num_workers
ds = FaceMask(data_root, cropsize=cropsize, mode='train')
dl = DataLoader(ds,
batch_size=n_img_all_gpu,
shuffle=True,
num_workers=num_workers,
drop_last=True)
ds_eval = FaceMask(data_root, cropsize=cropsize, mode='val')
dl_eval = DataLoader(ds_eval,
batch_size=n_img_all_gpu,
shuffle=True,
num_workers=num_workers,
drop_last=True)
ignore_idx = -100
net = BiSeNet(n_classes=n_classes)
# load last weights
if opt.load_weights is not None:
if opt.load_weights.endswith('.pth'):
weights_path = opt.load_weights
else:
weights_path = get_last_weights(opt.saved_path)
try:
last_step = int(
os.path.basename(weights_path).split('_')[-1].split('.')[0])
except:
last_step = 0
try:
ret = net.load_state_dict(torch.load(weights_path), strict=False)
except RuntimeError as e:
print(f'[Warning] Ignoring {e}')
print(
'[Warning] Don\'t panic if you see this, this might be because you load a pretrained weights '
'with different number of classes. The rest of the weights should be loaded already.'
)
print(
f'[Info] loaded weights: {os.path.basename(weights_path)}, resuming checkpoint from step: {last_step}'
)
else:
last_step = 0
print('[Info] initializing weights...')
writer = SummaryWriter(
opt.log_path +
f'/{datetime.datetime.now().strftime("%Y%m%d-%H%M%S")}/')
net = net.cuda()
net = nn.DataParallel(net)
score_thres = 0.7
n_min = n_img_all_gpu * cropsize[0] * cropsize[1] // opt.batch_size
LossP = OhemCELoss(thresh=score_thres, n_min=n_min, ignore_lb=ignore_idx)
Loss2 = OhemCELoss(thresh=score_thres, n_min=n_min, ignore_lb=ignore_idx)
Loss3 = OhemCELoss(thresh=score_thres, n_min=n_min, ignore_lb=ignore_idx)
# optimizer
momentum = 0.9
weight_decay = 5e-4
lr_start = opt.lr
max_iter = 80000
power = 0.9
warmup_steps = 1000
warmup_start_lr = 1e-5
optim = Optimizer(model=net.module,
lr0=lr_start,
momentum=momentum,
wd=weight_decay,
warmup_steps=warmup_steps,
warmup_start_lr=warmup_start_lr,
max_iter=max_iter,
power=power)
scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optim.optim,
patience=3,
verbose=True)
# train loop
loss_avg = []
step = max(0, last_step)
max_iter = len(dl)
best_epoch = 0
epoch = 0
best_loss = 1e5
net.train()
try:
for epoch in range(opt.num_epochs):
last_epoch = step // max_iter
if epoch < last_epoch:
continue
epoch_loss = []
progress_bar = tqdm(dl)
for iter, data in enumerate(progress_bar):
if iter < step - last_epoch * max_iter:
progress_bar.update()
continue
try:
im = data['img']
lb = data['label']
lb = torch.squeeze(lb, 1)
im = im.cuda()
lb = lb.cuda()
optim.zero_grad()
out, out16, out32 = net(im)
lossp = LossP(out, lb)
loss2 = Loss2(out16, lb)
loss3 = Loss3(out32, lb)
loss = lossp + loss2 + loss3
if loss == 0 or not torch.isfinite(loss):
continue
loss.backward()
optim.step()
loss_avg.append(loss.item())
# print training log message
# progress_bar.set_description(
# 'Epoch: {}/{}. Iteration: {}/{}. p_loss: {:.5f}. 2_loss: {:.5f}. 3_loss: {:.5f}. loss_avg: {:.5f}'.format(
# epoch, opt.num_epochs, iter + 1, max_iter, lossp.item(),
# loss2.item(), loss3.item(), loss.item()))
print(
'p_loss: {:.5f}. 2_loss: {:.5f}. 3_loss: {:.5f}. loss_avg: {:.5f}'
.format(lossp.item(), loss2.item(), loss3.item(),
loss.item()))
writer.add_scalars('Lossp', {'train': lossp}, step)
writer.add_scalars('loss2', {'train': loss2}, step)
writer.add_scalars('loss3', {'train': loss3}, step)
writer.add_scalars('loss_avg', {'train': loss}, step)
# log learning_rate
lr = optim.lr
writer.add_scalar('learning_rate', lr, step)
step += 1
if step % opt.save_interval == 0 and step > 0:
save_checkpoint(net, f'Bisenet_{epoch}_{step}.pth')
print('checkpoint...')
except Exception as e:
print('[Erro]', traceback.format_exc())
print(e)
continue
scheduler.step(np.mean(epoch_loss))
if epoch % opt.val_interval == 0:
net.eval()
loss_p = []
loss_2 = []
loss_3 = []
for iter, data in enumerate(dl_eval):
with torch.no_grad():
im = data['img']
lb = data['label']
lb = torch.squeeze(lb, 1)
im = im.cuda()
lb = lb.cuda()
out, out16, out32 = net(im)
lossp = LossP(out, lb)
loss2 = Loss2(out16, lb)
loss3 = Loss3(out32, lb)
loss = lossp + loss2 + loss3
if loss == 0 or not torch.isfinite(loss):
continue
loss_p.append(lossp.item())
loss_2.append(loss2.item())
loss_3.append(loss3.item())
lossp = np.mean(loss_p)
loss2 = np.mean(loss_2)
loss3 = np.mean(loss_3)
loss = lossp + loss2 + loss3
print(
'Val. Epoch: {}/{}. p_loss: {:1.5f}. 2_loss: {:1.5f}. 3_loss: {:1.5f}. Total_loss: {:1.5f}'
.format(epoch, opt.num_epochs, lossp, loss2, loss3, loss))
writer.add_scalars('Total_loss', {'val': loss}, step)
writer.add_scalars('p_loss', {'val': lossp}, step)
writer.add_scalars('2_loss', {'val': loss2}, step)
writer.add_scalars('3_loss', {'val': loss3}, step)
if loss + opt.es_min_delta < best_loss:
best_loss = loss
best_epoch = epoch
save_checkpoint(net, f'Bisenet_{epoch}_{step}.pth')
net.train() # ??
# Early stopping
if epoch - best_epoch > opt.es_patience > 0:
print(
'[Info] Stop training at epoch {}. The lowest loss achieved is {}'
.format(epoch, loss))
break
except KeyboardInterrupt:
save_checkpoint(net, f'Bisenet_{epoch}_{step}.pth')
writer.close()
writer.close()
def cluster(args, outpath):
# network setup
n_classes = 19
net = BiSeNet(n_classes=n_classes).cuda()
save_pth = 'resources/face_parsing_pytorch/res/cp/79999_iter.pth'
net.load_state_dict(torch.load(save_pth))
net.eval()
to_tensor = transforms.Compose([
transforms.Resize((512, 512), Image.BILINEAR),
transforms.ToTensor(),
transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),
])
subsets = ['reals_easiest', 'fakes_easiest']
PatchInfo = namedtuple('PatchInfo', ['patch', 'pos', 'file', 'value'])
for subset in subsets:
print(subset)
path = args.path
cluster_dir = os.path.join(outpath, subset + '_clusters')
os.makedirs(cluster_dir, exist_ok =True)
with open(os.path.join(path, subset+'_files.txt')) as f:
files = [line.strip() for line in f]
patches = np.load(os.path.join(path, subset + '.npz'))
which_model_netD = patches['which_model_netD']
fineSize = patches['finesize']
rfs = rfutil.find_rf_patches(which_model_netD, fineSize)
clusters = defaultdict(list)
clusters_baseline = Counter() # a counter for segmentations of random patches
# assign each patch to a cluster based on segmentation
for index, (patch, pos, value, file) in tqdm(enumerate(
zip(patches['patch'], patches['pos'], patches['value'], files)), total=len(files)):
image = Image.open(file).convert('RGB')
with torch.no_grad():
tensor = to_tensor(image)[None].cuda()
out = net(tensor)[0]
parsing = out[0].cpu().numpy().argmax(0)
interp = cv2.resize(parsing, (fineSize, fineSize), interpolation=cv2.INTER_NEAREST)[None]
seg_patch = rfutil.get_patch_from_img(interp, rfs[(pos[0], pos[1])], patches['rf'], pad_value=-1)
# how many pixels of each segmentation class in the patch
counter_patch = Counter(seg_patch[0].ravel())
# how many pixels of each segmentation class in the full img
counter_full = Counter(interp[0].ravel())
# normalize: omit padding value from normalization
counter_norm = {k: counter_patch[k] / counter_full[k] for k in counter_patch.keys() if k != -1}
cluster_id = max(counter_norm, key=counter_norm.get)
cluster_label = cluster_assn[cluster_id]
clusters[cluster_label].append(PatchInfo(patch, pos, file, value))
# pick a random patch for baseline
random_patch = random.choice(list(rfs.values()))
seg_patch = rfutil.get_patch_from_img(interp, random_patch, patches['rf'], pad_value=-1)
counter_patch = Counter(seg_patch[0].ravel())
counter_full = Counter(interp[0].ravel())
# omit padding value from normalization
counter_norm = {k: counter_patch[k] / counter_full[k] for k in counter_patch.keys() if k != -1}
cluster_id = max(counter_norm, key=counter_norm.get)
cluster_label = cluster_assn[cluster_id]
clusters_baseline[cluster_label] += 1
# plot each cluster in a grid
counts, labels = [], []
infos = []
for index, (k,v) in enumerate(sorted(clusters.items(), key=lambda item: len(item[1]))[::-1]):
line = '%d: %s, %d patches' % (index, k, len(v))
print(line)
infos.append(line)
counts.append(len(v))
labels.append(k)
cluster = np.asarray([patchinfo.patch for patchinfo in v])
files = [patchinfo.file for patchinfo in clusters[k]]
normalized = (cluster[-225:] * 0.5) + 0.5 # at most 15x15 grid
grid = imutil.imgrid(np.uint8(normalized * 255), pad=0, cols=int(np.ceil(np.sqrt(normalized.shape[0]))))
grid_im = Image.fromarray(grid)
grid_im.save(os.path.join(cluster_dir, 'cluster_%d.png' % index))
np.savez(os.path.join(cluster_dir, 'cluster_%d.npz' % index),
patch=cluster, rf=patches['rf'],
finesize=patches['finesize'],
outsize=patches['outsize'],
which_model_netD=patches['which_model_netD'],
pos=np.array([patchinfo.pos for patchinfo in v]),
value=np.array([patchinfo.value for patchinfo in v]))
with open(os.path.join(cluster_dir, 'cluster_%d.txt' % index), 'w') as f:
[f.write('%s\n' % file) for file in files]
# histogram
f, ax = plt.subplots(1, 1, figsize=(6, 4))
ax.bar(range(1, len(labels) + 1), counts)
ax.set_xticks(range(1, len(labels) + 1))
ax.set_xticklabels(labels, rotation='vertical')
ax.set_ylabel('count')
f.savefig(os.path.join(cluster_dir, 'histogram.pdf'),
bbox_inches='tight')
# write counts to file
with open(os.path.join(cluster_dir, 'counts.txt'), 'w') as f:
[f.write('%s\n' % line) for line in infos]
# write random patch baseline to file
infos = []
for index, (k,v) in enumerate(sorted(clusters_baseline.items(), key=lambda item: item[1])[::-1]):
infos.append('%d: %s, %d patches' % (index, k, v))
with open(os.path.join(cluster_dir, 'baseline.txt'), 'w') as f:
[f.write('%s\n' % line) for line in infos]
def evaluate(respth='./res/test_res',
dspth='./data',
cp='model_final_diss.pth'):
if not os.path.exists(respth):
os.makedirs(respth)
n_classes = 19
net = BiSeNet(n_classes=n_classes)
net.cuda()
save_pth = osp.join(respth, 'cp', cp)
net.load_state_dict(torch.load(save_pth))
net.eval()
no_iter = str(int(cp.split('_')[0]))
org_respth = respth[:]
respth = os.path.join(respth, no_iter)
if not os.path.exists(respth):
os.makedirs(respth)
to_tensor = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),
])
''' added '''
cropsize = [448, 448]
n_img_per_gpu = 16
data_root = '/home/jihyun/workspace/face_parsing/dataset/CelebAMask-HQ/'
ds = FaceMask(data_root, cropsize=cropsize, mode='val')
dl = DataLoader(ds, batch_size=16, shuffle=False, drop_last=True)
n_min = n_img_per_gpu * cropsize[0] * cropsize[1] // 16
score_thres = 0.7
ignore_idx = -100
loss_avg = []
LossP = OhemCELoss(thresh=score_thres, n_min=n_min, ignore_lb=ignore_idx)
Loss2 = OhemCELoss(thresh=score_thres, n_min=n_min, ignore_lb=ignore_idx)
Loss3 = OhemCELoss(thresh=score_thres, n_min=n_min, ignore_lb=ignore_idx)
with torch.no_grad():
for i, sample in enumerate(dl):
im, lb = sample
im = im.cuda()
lb = lb.cuda()
lb = torch.squeeze(lb, 1)
out, out16, out32 = net(im)
lossp = LossP(out, lb)
loss2 = Loss2(out16, lb)
loss3 = Loss3(out32, lb)
loss = lossp + loss2 + loss3
loss_avg.append(loss.item())
loss_avg = sum(loss_avg) / len(loss_avg)
f = open(osp.join(org_respth, 'loss.log'), 'a')
f.write(' eval_loss: ' + str(loss_avg) + '\n')
f.close()
for image_path in os.listdir(dspth):
img = Image.open(osp.join(dspth, image_path))
image = img.resize((512, 512), Image.BILINEAR)
img = to_tensor(image)
img = torch.unsqueeze(img, 0)
img = img.cuda()
out, out16, out32 = net(img)
parsing = out.squeeze(0).cpu().numpy().argmax(0)
vis_parsing_maps(image,
parsing,
stride=1,
save_im=True,
save_path=osp.join(respth, image_path))
def allowed_file(filename):
return '.' in filename and filename.rsplit(
'.', 1)[1].lower() in ['png', 'jpeg', 'jpg']
to_tensor = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),
])
# load model
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
NCLASSES = 19
model = BiSeNet(n_classes=NCLASSES)
model.load_state_dict(torch.load(model_src, map_location=torch.device(device)))
model.eval()
@app.route('/')
def home():
return render_template("home.html")
@app.route('/predict', methods=['POST'])
def predict():
# get file
img_file = request.files['img']
if img_file and allowed_file(img_file.filename):
npimg = np.fromstring(img_file.read(), np.uint8)
img = cv2.imdecode(npimg, cv2.IMREAD_COLOR)
def setup(opts):
net = BiSeNet(n_classes=19)
net.cuda()
net.load_state_dict(torch.load(opts['checkpoint']))
net.eval()
return net
def evaluate(source='./data', dest='./res/test_res', cp='evaluate.pth'):
n_classes = 19
net = BiSeNet(n_classes=n_classes)
net.cuda()
save_pth = osp.join('res/cp', cp)
net.load_state_dict(torch.load(save_pth))
net.eval()
to_tensor = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),
])
if path.isfile(source):
full_path = True
#dspth = np.sort(np.loadtxt(source, dtype=np.str))
dspth = np.sort(
np.asarray(pd.read_csv(source, delimiter=' ',
header=None)).flatten())
else:
full_path = False
dspth = listdir(source)
with torch.no_grad():
for image_name in tqdm(dspth):
if not full_path:
image_path = path.join(source, image_name)
else:
image_path = image_name
image_name = path.split(image_name)[1]
img = Image.open(image_path)
if full_path:
sub_folder = path.basename(
path.normpath(path.split(image_path)[0]))
dest_path = path.join(dest, sub_folder)
if not path.exists(dest_path):
makedirs(dest_path)
else:
dest_path = dest
save_path = osp.join(dest_path, image_name)
if os.path.isfile(save_path[:-4] + '_mask.png'):
print('Skipping...')
continue
image = img.resize((512, 512), Image.BILINEAR)
img = to_tensor(image)
img = torch.unsqueeze(img, 0)
img = img.cuda()
out = net(img)[0]
parsing = out.squeeze(0).cpu().numpy().argmax(0)
# print(parsing)
# print(np.unique(parsing))
vis_parsing_maps(image,
parsing,
stride=1,
save_im=True,
save_path=save_path)
