def hed_edges(image):
import cv2 as cv
# based on https://github.com/opencv/opencv/blob/master/samples/dnn/edge_detection.py
class CropLayer(object):
def __init__(self, params, blobs):
self.xstart = 0
self.xend = 0
self.ystart = 0
self.yend = 0
# Our layer receives two inputs. We need to crop the first input blob
# to match a shape of the second one (keeping batch size and number of channels)
def getMemoryShapes(self, inputs):
inputShape, targetShape = inputs[0], inputs[1]
batchSize, numChannels = inputShape[0], inputShape[1]
height, width = targetShape[2], targetShape[3]
self.ystart = int((inputShape[2] - targetShape[2]) / 2)
self.xstart = int((inputShape[3] - targetShape[3]) / 2)
self.yend = self.ystart + height
self.xend = self.xstart + width
return [[batchSize, numChannels, height, width]]
def forward(self, inputs):
return [
inputs[0][:, :, self.ystart:self.yend, self.xstart:self.xend]
]
# Load the pretrained model (source: https://github.com/s9xie/hed)
script_path = Path(__file__).parent.absolute()
hed_path = Path.joinpath(script_path, 'HED')
net = cv.dnn.readNetFromCaffe(
str(hed_path / 'deploy.prototxt'),
str(hed_path / 'hed_pretrained_bsds.caffemodel'))
cv.dnn_registerLayer('Crop', CropLayer)
image = cv.resize(image, (image.shape[1], image.shape[0]))
# prepare image as input dataset (mean values from full image dataset)
inp = cv.dnn.blobFromImage(
image,
scalefactor=1.0,
size=(image.shape[1], image.shape[0]), #w,h
mean=(104.00698793, 116.66876762, 122.67891434),
swapRB=False,
crop=False)
net.setInput(inp)
out = net.forward()
cv.dnn_unregisterLayer('Crop') # get rid of issues when run in a loop
out = out[0, 0]
return out
def holistically_nested(image):
p_i("Starting Holistically-Nested Edge Detection...")
net = cv2.dnn.readNetFromCaffe(
"data/hed_model/deploy.prototxt",
"data/hed_model/" + "hed_pretrained_bsds.caffemodel",
)
height, width = image.shape[:2]
cv2.dnn_registerLayer("Crop", CropLayer)
blob = cv2.dnn.blobFromImage(image, size=(width, height))
net.setInput(blob)
hed = (255 * cv2.resize(net.forward()[0, 0],
(width, height))).astype("uint8")
p_i("Holistically-Nested Edge Detection complete!")
cv2.dnn_unregisterLayer("Crop")
return hed
def hedConvert(image):
protoPath = os.path.sep.join(["hed_model", "deploy.prototxt"])
modelPath = os.path.sep.join(
["hed_model", "hed_pretrained_bsds.caffemodel"])
net = cv2.dnn.readNetFromCaffe(protoPath, modelPath)
cv2.dnn_registerLayer("Crop", CropLayer)
(H, W) = image.shape[:2]
blob = cv2.dnn.blobFromImage(image,
scalefactor=1.0,
size=(W, H),
mean=(104.00698793, 116.66876762,
122.67891434),
swapRB=False,
crop=False)
net.setInput(blob)
hed = net.forward()
hed = cv2.resize(hed[0, 0], (W, H))
hed = (255 * hed).astype("uint8")
cv2.dnn_unregisterLayer("Crop")
return hed
def test_custom_layer(self):
class CropLayer(object):
def __init__(self, params, blobs):
self.xstart = 0
self.xend = 0
self.ystart = 0
self.yend = 0
# Our layer receives two inputs. We need to crop the first input blob
# to match a shape of the second one (keeping batch size and number of channels)
def getMemoryShapes(self, inputs):
inputShape, targetShape = inputs[0], inputs[1]
batchSize, numChannels = inputShape[0], inputShape[1]
height, width = targetShape[2], targetShape[3]
self.ystart = (inputShape[2] - targetShape[2]) // 2
self.xstart = (inputShape[3] - targetShape[3]) // 2
self.yend = self.ystart + height
self.xend = self.xstart + width
return [[batchSize, numChannels, height, width]]
def forward(self, inputs):
return [
inputs[0][:, :, self.ystart:self.yend,
self.xstart:self.xend]
]
cv.dnn_registerLayer('CropCaffe', CropLayer)
proto = '''
name: "TestCrop"
input: "input"
input_shape
{
dim: 1
dim: 2
dim: 5
dim: 5
}
input: "roi"
input_shape
{
dim: 1
dim: 2
dim: 3
dim: 3
}
layer {
name: "Crop"
type: "CropCaffe"
bottom: "input"
bottom: "roi"
top: "Crop"
}'''
net = cv.dnn.readNetFromCaffe(bytearray(proto.encode()))
for backend, target in self.dnnBackendsAndTargets:
if backend != cv.dnn.DNN_BACKEND_OPENCV:
continue
printParams(backend, target)
net.setPreferableBackend(backend)
net.setPreferableTarget(target)
src_shape = [1, 2, 5, 5]
dst_shape = [1, 2, 3, 3]
inp = np.arange(0, np.prod(src_shape),
dtype=np.float32).reshape(src_shape)
roi = np.empty(dst_shape, dtype=np.float32)
net.setInput(inp, "input")
net.setInput(roi, "roi")
out = net.forward()
ref = inp[:, :, 1:4, 1:4]
normAssert(self, out, ref)
cv.dnn_unregisterLayer('CropCaffe')
out_h = inp_w * ratio
start = int(center_h - out_h // 2)
end = int(center_h + out_h // 2)
person_img = person_img[start:end, ...]
else:
center_w = inp_w // 2
out_w = inp_h / ratio
start = int(center_w - out_w // 2)
end = int(center_w + out_w // 2)
person_img = person_img[:, start:end, :]
cloth_img = cv.imread(args.input_cloth)
pose = get_pose_map(person_img, findFile(args.openpose_proto),
findFile(args.openpose_model), args.backend, args.target)
segm_image = parse_human(person_img, args.segmentation_model)
segm_image = cv.resize(segm_image, (192, 256), cv.INTER_LINEAR)
cv.dnn_registerLayer('Correlation', CorrelationLayer)
model = CpVton(args.gmm_model, args.tom_model, args.backend, args.target)
agnostic = model.prepare_agnostic(segm_image, person_img, pose)
warped_cloth = model.get_warped_cloth(cloth_img, agnostic)
output = model.get_tryon(agnostic, warped_cloth)
cv.dnn_unregisterLayer('Correlation')
winName = 'Virtual Try-On'
cv.namedWindow(winName, cv.WINDOW_AUTOSIZE)
cv.imshow(winName, output)
cv.waitKey()