def _build_matting_model(self):
class Args:
encoder = 'resnet50_GN_WS'
decoder = 'fba_decoder'
weights = os.path.join(os.path.dirname(__file__),
'FBA_Matting/FBA.pth')
args = Args()
self.model = build_model(args)
def perform_matting(image, trimap):
class Args:
encoder = 'resnet50_GN_WS'
decoder = 'fba_decoder'
weights = 'FBA.pth'
args = Args()
model = build_model(args)
model.eval()
return pred(image, trimap, model)
def __init__(self, model_dir, model_config):
parser = argparse.ArgumentParser()
# Model related arguments
class Args:
encoder = 'resnet50_GN_WS'
decoder = 'fba_decoder'
weights = 'FBA.pth'
args = Args()
model = build_model(args)
model.eval()
fg, bg, alpha = pred(image_np, trimap_np, model)
from networks.models import build_model
class Args:
def __init__(self):
self.encoder = 'resnet50_GN_WS'
self.decoder = 'fba_decoder'
self.weights = 'FBA.pth'
args = Args()
Model = build_model(args)
import io
app = Flask(__name__)
def get_array(arg):
return np.array(arg).astype('uint8')
class Matting_Args:
def __init__(self):
self.encoder = 'resnet50_GN_WS'
self.decoder = 'fba_decoder'
self.weights = '../models/FBA.pth'
args = Matting_Args()
matting_model = build_model(args)
matting_model.eval();
def get_response(new_bg,data):
image = get_array(data.get('image'))
response = requests.post('http://127.0.0.1:3000/',json = data)
if response.status_code == 406:
return jsonify({'output':image.tolist()})
h,w,_ = image.shape
trimap = get_array(response.json()['trimap'])
fg, bg, alpha = pred(image/255.0,trimap,matting_model)
combined = ((alpha[...,None]*image)).astype('uint8') + ((1-alpha)[...,None]*cv2.resize(new_bg,(w,h))).astype('uint8')
return jsonify({'output':combined.tolist()})
@app.route('/with_bg',methods=['POST'])
def extraction():
output[0].cpu().numpy().transpose((1, 2, 0)), (w, h), cv2.INTER_LANCZOS4,
)
alpha = output[:, :, 0]
fg = output[:, :, 1:4]
bg = output[:, :, 4:7]
alpha[trimap_np[:, :, 0] == 1] = 0
alpha[trimap_np[:, :, 1] == 1] = 1
fg[alpha == 1] = image_np[alpha == 1]
bg[alpha == 0] = image_np[alpha == 0]
return fg, bg, alpha
if __name__ == "__main__":
parser = argparse.ArgumentParser()
# Model related arguments
parser.add_argument("--encoder", default="resnet50_GN_WS", help="Encoder model")
parser.add_argument("--decoder", default="fba_decoder", help="Decoder model")
parser.add_argument("--weights", default="FBA.pth")
parser.add_argument("--image_dir", default="./examples/images", help="")
parser.add_argument(
"--trimap_dir", default="./examples/trimaps", help="",
)
parser.add_argument("--output_dir", default="./examples/predictions", help="")
parser.add_argument("--device", default="cpu", help="Device for inference on")
args = parser.parse_args()
model = build_model(args).to(args.device)
model.eval()
predict_fba_folder(model, args)
