mirror of
https://github.com/deepinsight/insightface.git
synced 2026-05-16 13:46:15 +00:00
67 lines
2.4 KiB
Python
67 lines
2.4 KiB
Python
#!/usr/bin/env python
|
|
# -*- encoding: utf-8 -*-
|
|
"""
|
|
@Author : Qingping Zheng
|
|
@Contact : qingpingzheng2014@gmail.com
|
|
@File : edges.py
|
|
@Time : 10/01/21 00:00 PM
|
|
@Desc :
|
|
@License : Licensed under the Apache License, Version 2.0 (the "License");
|
|
@Copyright : Copyright 2022 The Authors. All Rights Reserved.
|
|
"""
|
|
from __future__ import absolute_import
|
|
from __future__ import division
|
|
from __future__ import print_function
|
|
|
|
import torch
|
|
import torch.nn.functional as F
|
|
import torch.nn as nn
|
|
|
|
from inplace_abn import InPlaceABNSync
|
|
|
|
|
|
class Edges(nn.Module):
|
|
|
|
def __init__(self, abn=InPlaceABNSync, in_fea=[256,512,1024], mid_fea=256, out_fea=2):
|
|
super(Edges, self).__init__()
|
|
|
|
self.conv1 = nn.Sequential(
|
|
nn.Conv2d(in_fea[0], mid_fea, kernel_size=1, padding=0, dilation=1, bias=False),
|
|
abn(mid_fea)
|
|
)
|
|
self.conv2 = nn.Sequential(
|
|
nn.Conv2d(in_fea[1], mid_fea, kernel_size=1, padding=0, dilation=1, bias=False),
|
|
abn(mid_fea)
|
|
)
|
|
self.conv3 = nn.Sequential(
|
|
nn.Conv2d(in_fea[2], mid_fea, kernel_size=1, padding=0, dilation=1, bias=False),
|
|
abn(mid_fea)
|
|
)
|
|
self.conv4 = nn.Conv2d(mid_fea,out_fea, kernel_size=3, padding=1, dilation=1, bias=True)
|
|
self.conv5_b = nn.Conv2d(out_fea*3,2, kernel_size=1, padding=0, dilation=1, bias=True)
|
|
self.conv5 = nn.Conv2d(out_fea*3,out_fea, kernel_size=1, padding=0, dilation=1, bias=True)
|
|
|
|
|
|
def forward(self, x1, x2, x3):
|
|
_, _, h, w = x1.size()
|
|
|
|
edge1_fea = self.conv1(x1)
|
|
edge1 = self.conv4(edge1_fea)
|
|
edge2_fea = self.conv2(x2)
|
|
edge2 = self.conv4(edge2_fea)
|
|
edge3_fea = self.conv3(x3)
|
|
edge3 = self.conv4(edge3_fea)
|
|
|
|
edge2_fea = F.interpolate(edge2_fea, size=(h, w), mode='bilinear',align_corners=True)
|
|
edge3_fea = F.interpolate(edge3_fea, size=(h, w), mode='bilinear',align_corners=True)
|
|
edge2 = F.interpolate(edge2, size=(h, w), mode='bilinear',align_corners=True)
|
|
edge3 = F.interpolate(edge3, size=(h, w), mode='bilinear',align_corners=True)
|
|
|
|
edge = torch.cat([edge1, edge2, edge3], dim=1)
|
|
edge_fea = torch.cat([edge1_fea, edge2_fea, edge3_fea], dim=1)
|
|
semantic_edge = self.conv5(edge)
|
|
binary_edge = self.conv5_b(edge)
|
|
|
|
return binary_edge, semantic_edge, edge_fea
|
|
|