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(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='./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(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(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(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))
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(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 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 __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 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 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 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(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='')
def setup(opts):
net = BiSeNet(n_classes=19)
net.cuda()
net.load_state_dict(torch.load(opts['checkpoint']))
net.eval()
return net
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)))
total_f_score = []
class_f_score = [[] for i in range(n_classes)]
with torch.no_grad():
for i, sample in enumerate(dl):
if i == 5: break
if (i + 1) % 5 == 0:
print('processing {}-th batch...'.format(str(i + 1)))
im, lb = sample
im = im.cuda()
lb = lb.cuda()
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))
#!/usr/bin/python
# -*- encoding: utf-8 -*-
import torch
import os
from model import BiSeNet
import os.path as osp
import numpy as np
from PIL import Image
import torchvision.transforms as transforms
import cv2
n_classes = 19
net = BiSeNet(n_classes=n_classes)
net.cuda()
net.load_state_dict(torch.load('cp/79999_iter.pth'))
net.eval()
def vis_parsing_maps(im,
parsing_anno,
stride,
save_im=False,
save_path='vis_results/parsing_map_on_im.jpg'):
# Colors for all 20 parts
part_colors = [[255, 0, 0], [255, 85, 0], [255, 170, 0], [255, 0, 85],
[255, 0, 170], [0, 255, 0], [85, 255, 0], [170, 255, 0],
[0, 255, 85], [0, 255, 170], [0, 0, 255], [85, 0, 255],
[170, 0, 255], [0, 85, 255], [0, 170, 255], [255, 255, 0],
[255, 255, 85], [255, 255, 170], [255, 0, 255],
[255, 85, 255], [255, 170, 255], [0, 255, 255],
def train():
args = parse_args()
torch.cuda.set_device(args.local_rank)
dist.init_process_group(backend='nccl',
init_method='tcp://127.0.0.1:33271',
world_size=torch.cuda.device_count(),
rank=args.local_rank)
setup_logger(respth)
## dataset
n_classes = 19
n_img_per_gpu = 8
n_workers = 4
cropsize = [1024, 1024]
ds = CityScapes('./data', cropsize=cropsize, mode='train')
sampler = torch.utils.data.distributed.DistributedSampler(ds)
dl = DataLoader(ds,
batch_size=n_img_per_gpu,
shuffle=False,
sampler=sampler,
num_workers=n_workers,
pin_memory=True,
drop_last=True)
## model
ignore_idx = 255
net = BiSeNet(n_classes=n_classes)
if not args.ckpt is None:
net.load_state_dict(torch.load(args.ckpt, map_location='cpu'))
net.cuda()
net.train()
net = nn.parallel.DistributedDataParallel(net,
device_ids=[
args.local_rank,
],
output_device=args.local_rank)
score_thres = 0.7
n_min = n_img_per_gpu * cropsize[0] * cropsize[1] // 16
criteria_p = OhemCELoss(thresh=score_thres,
n_min=n_min,
ignore_lb=ignore_idx)
criteria_16 = OhemCELoss(thresh=score_thres,
n_min=n_min,
ignore_lb=ignore_idx)
criteria_32 = OhemCELoss(thresh=score_thres,
n_min=n_min,
ignore_lb=ignore_idx)
## optimizer
momentum = 0.9
weight_decay = 5e-4
lr_start = 1e-2
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)
## train loop
msg_iter = 50
loss_avg = []
st = glob_st = time.time()
diter = iter(dl)
epoch = 0
for it in range(max_iter):
try:
im, lb = next(diter)
if not im.size()[0] == n_img_per_gpu: raise StopIteration
except StopIteration:
epoch += 1
sampler.set_epoch(epoch)
diter = iter(dl)
im, lb = next(diter)
im = im.cuda()
lb = lb.cuda()
H, W = im.size()[2:]
lb = torch.squeeze(lb, 1)
optim.zero_grad()
out, out16, out32 = net(im)
lossp = criteria_p(out, lb)
loss2 = criteria_16(out16, lb)
loss3 = criteria_32(out32, lb)
loss = lossp + loss2 + loss3
loss.backward()
optim.step()
loss_avg.append(loss.item())
## print training log message
if (it + 1) % msg_iter == 0:
loss_avg = sum(loss_avg) / len(loss_avg)
lr = optim.lr
ed = time.time()
t_intv, glob_t_intv = ed - st, ed - glob_st
eta = int((max_iter - it) * (glob_t_intv / it))
eta = str(datetime.timedelta(seconds=eta))
msg = ', '.join([
'it: {it}/{max_it}',
'lr: {lr:4f}',
'loss: {loss:.4f}',
'eta: {eta}',
'time: {time:.4f}',
]).format(it=it + 1,
max_it=max_iter,
lr=lr,
loss=loss_avg,
time=t_intv,
eta=eta)
logger.info(msg)
loss_avg = []
st = ed
## dump the final model
save_pth = osp.join(respth, 'model_final.pth')
net.cpu()
state = net.module.state_dict() if hasattr(net,
'module') else net.state_dict()
if dist.get_rank() == 0: torch.save(state, save_pth)
logger.info('training done, model saved to: {}'.format(save_pth))
def train(fintune_model, data_root, respth):
# dataset
n_classes = 19
n_img_per_gpu = 16
n_workers = 8
cropsize = [448, 448]
ds = FaceMask(data_root, cropsize=cropsize, mode='train')
# sampler = torch.utils.data.distributed.DistributedSampler(ds)
dl = DataLoader(ds,
batch_size=n_img_per_gpu,
shuffle=True,
num_workers=n_workers,
pin_memory=True,
drop_last=True)
# model
ignore_idx = -100
use_cuda = torch.cuda.is_available()
device = torch.device("cuda:0" if use_cuda else "cpu")
net = BiSeNet(n_classes=n_classes)
net = net.to(device)
if os.access(fintune_model, os.F_OK) and (fintune_model
is not None): # checkpoint
chkpt = torch.load(fintune_model, map_location=device)
net.load_state_dict(chkpt)
print('load fintune model : {}'.format(fintune_model))
else:
print('no fintune model')
score_thres = 0.7
n_min = n_img_per_gpu * cropsize[0] * cropsize[1] // 16
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 = 1e-2
max_epoch = 1000
optim = Optimizer.SGD(net.parameters(),
lr=lr_start,
momentum=momentum,
weight_decay=weight_decay)
## train loop
msg_iter = 50
loss_avg = []
st = glob_st = time.time()
# diter = iter(dl)
epoch = 0
flag_change_lr_cnt = 0 # 学习率更新计数器
init_lr = lr_start # 学习率
best_loss = np.inf
loss_mean = 0. # 损失均值
loss_idx = 0. # 损失计算计数器
print('start training ~')
it = 0
for epoch in range(max_epoch):
net.train()
# 学习率更新策略
if loss_mean != 0.:
if best_loss > (loss_mean / loss_idx):
flag_change_lr_cnt = 0
best_loss = (loss_mean / loss_idx)
else:
flag_change_lr_cnt += 1
if flag_change_lr_cnt > 30:
init_lr = init_lr * 0.1
set_learning_rate(optimizer, init_lr)
flag_change_lr_cnt = 0
loss_mean = 0. # 损失均值
loss_idx = 0. # 损失计算计数器
for i, (im, lb) in enumerate(dl):
im = im.cuda()
lb = lb.cuda()
H, W = im.size()[2:]
lb = torch.squeeze(lb, 1)
optim.zero_grad()
out, out16, out32 = net(im)
lossp = LossP(out, lb)
loss2 = Loss2(out16, lb)
loss3 = Loss3(out32, lb)
loss = lossp + loss2 + loss3
loss_mean += loss.item()
loss_idx += 1.
loss.backward()
optim.step()
if it % msg_iter == 0:
print('epoch <{}/{}> -->> <{}/{}> -> iter {} : loss {:.5f}, loss_mean :{:.5f}, best_loss :{:.5f},lr :{:.6f},batch_size : {}'.\
format(epoch,max_epoch,i,int(ds.__len__()/n_img_per_gpu),it,loss.item(),loss_mean/loss_idx,best_loss,init_lr,n_img_per_gpu))
# print(msg)
if (it) % 500 == 0:
state = net.module.state_dict() if hasattr(
net, 'module') else net.state_dict()
torch.save(state, respth + '/model/face_parse_latest.pth')
# evaluate(dspth='./images', cp='{}_iter.pth'.format(it))
it += 1
torch.save(state,
respth + '/model/face_parse_epoch_{}.pth'.format(epoch))
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)
import torchvision.transforms as transforms
import cv2
import os
from skimage.morphology import remove_small_holes, remove_small_objects
import sys
sys.path += os.path.abspath(__file__)
modelFile = "models/face_detector.pb"
configFile = "models/face_detector.pbtxt"
face_detector = cv2.dnn.readNetFromTensorflow(modelFile, configFile)
n_classes = 19
seg_net = BiSeNet(n_classes=n_classes)
seg_net.load_state_dict(torch.load('models/segmentation.pth', map_location='cpu'))
parser = argparse.ArgumentParser()
args = parser.parse_args()
args.resume = "models/matting.pth"
args.stage = 1
args.crop_or_resize = "whole"
args.max_size = 512
args.cuda = False
# init model
matting_model = net.VGG16(args)
ckpt = torch.load(args.resume, map_location='cpu')
matting_model.load_state_dict(ckpt['state_dict'], strict=True)
def segmentation(image):
class BISENET:
def __init__(self, model_path, csv_path):
# retrieve label info
self.label_info = get_label_info(csv_path)
# build model and load weight
self.model = BiSeNet(12, 'resnet18')
self.device = torch.device(
'cuda:0' if torch.cuda.is_available() else 'cpu')
self.model.load_state_dict(torch.load(model_path))
self.model.to(self.device).eval()
self.transform = transforms.Compose([
transforms.ToPILImage(),
transforms.ToTensor(),
transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225))
])
# retrieve person color
self.person_color = self.label_info['Pedestrian'][:-1]
def predict_on_image(self, image):
# transform image to tensor
image = self.transform(image[:, :, ::-1]).to(self.device)
# prediction map
predict = self.model(image.unsqueeze(0)).squeeze()
# encode to class index
predict = reverse_one_hot(predict).cpu().numpy()
# encode to color code
predict = colour_code_segmentation(predict,
self.label_info).astype(np.uint8)
# get bbox output
predict, bboxes, num_people = self.bbox_output(predict)
return predict, bboxes, num_people
def bbox_output(self, predict):
# get a binary mask with persons color white and background black
person_mask = np.zeros(predict.shape, dtype=np.uint8)
person_mask[np.all(predict == self.person_color,
axis=-1)] = [255, 255, 255]
person_mask = cv2.cvtColor(person_mask, cv2.COLOR_BGR2GRAY)
# label the mask image with connected-components algorithm
label_image = label(person_mask)
# find the bbox regions
regions = regionprops(label_image)
bboxes = []
num_people = [0]
i = 1
for props in regions:
if props.area > 100:
minr, minc, maxr, maxc = props.bbox
bboxes += [props.bbox]
num_people += [i]
predict = cv2.rectangle(predict, (minc, minr), (maxc, maxr),
(0, 255, 0), 2, cv2.LINE_AA)
predict = cv2.putText(predict, f'person{i}', (minc, minr - 10),
cv2.FONT_HERSHEY_SIMPLEX, 0.5,
(0, 255, 0), 2, cv2.LINE_AA)
i += 1
return predict, bboxes, num_people[-1]
