import argparse
from models.py_utils.data_parallel import DataParallel
import torch.nn as nn
from models.py_utils import TopPool, BottomPool, LeftPool, RightPool
from models.py_utils.utils import convolution, residual
from models.py_utils.utils import make_layer, make_layer_revr
from models.py_utils.kp_utils import _decode
from models.py_utils.kp_utils import make_tl_layer, make_br_layer, make_kp_layer, make_ct_layer
from models.py_utils.kp_utils import make_pool_layer, make_unpool_layer
from models.py_utils.kp_utils import make_merge_layer, make_inter_layer, make_cnv_layer
import os
import cv2
import json
import copy
import numpy as np
import torch
import matplotlib.pyplot as plt
from tqdm import tqdm
from config import system_configs
from utils import crop_image, normalize_
from external.nms import soft_nms, soft_nms_merge
colours = np.random.rand(80,3)
class kp_module(nn.Module):
def __init__(
self, n, dims, modules, layer=residual,
make_up_layer=make_layer, make_low_layer=make_layer,
make_hg_layer=make_layer, make_hg_layer_revr=make_layer_revr,
make_pool_layer=make_pool_layer, make_unpool_layer=make_unpool_layer,
make_merge_layer=make_merge_layer, **kwargs
):
super(kp_module, self).__init__()
self.n = n
curr_mod = modules[0]
next_mod = modules[1]
curr_dim = dims[0]
next_dim = dims[1]
self.up1 = make_up_layer(
3, curr_dim, curr_dim, curr_mod,
layer=layer, **kwargs
)
self.max1 = make_pool_layer(curr_dim)
self.low1 = make_hg_layer(
3, curr_dim, next_dim, curr_mod,
layer=layer, **kwargs
)
self.low2 = kp_module(
n - 1, dims[1:], modules[1:], layer=layer,
make_up_layer=make_up_layer,
make_low_layer=make_low_layer,
make_hg_layer=make_hg_layer,
make_hg_layer_revr=make_hg_layer_revr,
make_pool_layer=make_pool_layer,
make_unpool_layer=make_unpool_layer,
make_merge_layer=make_merge_layer,
**kwargs
) if self.n > 1 else \
make_low_layer(
3, next_dim, next_dim, next_mod,
layer=layer, **kwargs
)
self.low3 = make_hg_layer_revr(
3, next_dim, curr_dim, curr_mod,
layer=layer, **kwargs
)
self.up2 = make_unpool_layer(curr_dim)
self.merge = make_merge_layer(curr_dim)
def forward(self, x):
up1 = self.up1(x)
max1 = self.max1(x)
low1 = self.low1(max1)
low2 = self.low2(low1)
low3 = self.low3(low2)
up2 = self.up2(low3)
return self.merge(up1, up2)
class kp(nn.Module):
def __init__(
self, n, nstack, dims, modules, out_dim, pre=None, cnv_dim=256,
make_tl_layer=make_tl_layer, make_br_layer=make_br_layer, make_ct_layer=make_ct_layer,
make_cnv_layer=make_cnv_layer, make_heat_layer=make_kp_layer,
make_tag_layer=make_kp_layer, make_regr_layer=make_kp_layer,
make_up_layer=make_layer, make_low_layer=make_layer,
make_hg_layer=make_layer, make_hg_layer_revr=make_layer_revr,
make_pool_layer=make_pool_layer, make_unpool_layer=make_unpool_layer,
make_merge_layer=make_merge_layer, make_inter_layer=make_inter_layer,
kp_layer=residual
):
super(kp, self).__init__()
self.nstack = nstack
self._decode = _decode
self.K = 100
self.ae_threshold = 0.5
self.kernel = 3
self.input_size = 511
self.output_size = 128
curr_dim = dims[0]
self.pre = nn.Sequential(
convolution(7, 3, 128, stride=2),
residual(3, 128, 256, stride=2)
) if pre is None else pre
self.kps = nn.ModuleList([
kp_module(
n, dims, modules, layer=kp_layer,
make_up_layer=make_up_layer,
make_low_layer=make_low_layer,
make_hg_layer=make_hg_layer,
make_hg_layer_revr=make_hg_layer_revr,
make_pool_layer=make_pool_layer,
make_unpool_layer=make_unpool_layer,
make_merge_layer=make_merge_layer
) for _ in range(nstack)
])
self.cnvs = nn.ModuleList([
make_cnv_layer(curr_dim, cnv_dim) for _ in range(nstack)
])
self.tl_cnvs = nn.ModuleList([
make_tl_layer(cnv_dim) for _ in range(nstack)
])
self.br_cnvs = nn.ModuleList([
make_br_layer(cnv_dim) for _ in range(nstack)
])
self.ct_cnvs = nn.ModuleList([
make_ct_layer(cnv_dim) for _ in range(nstack)
])
## keypoint heatmaps
self.tl_heats = nn.ModuleList([
make_heat_layer(cnv_dim, curr_dim, out_dim) for _ in range(nstack)
])
self.br_heats = nn.ModuleList([
make_heat_layer(cnv_dim, curr_dim, out_dim) for _ in range(nstack)
])
self.ct_heats = nn.ModuleList([
make_heat_layer(cnv_dim, curr_dim, out_dim) for _ in range(nstack)
])
## tags
self.tl_tags = nn.ModuleList([
make_tag_layer(cnv_dim, curr_dim, 1) for _ in range(nstack)
])
self.br_tags = nn.ModuleList([
make_tag_layer(cnv_dim, curr_dim, 1) for _ in range(nstack)
])
for tl_heat, br_heat, ct_heat in zip(self.tl_heats, self.br_heats, self.ct_heats):
tl_heat[-1].bias.data.fill_(-2.19)
br_heat[-1].bias.data.fill_(-2.19)
ct_heat[-1].bias.data.fill_(-2.19)
self.inters = nn.ModuleList([
make_inter_layer(curr_dim) for _ in range(nstack - 1)
])
self.inters_ = nn.ModuleList([
nn.Sequential(
nn.Conv2d(curr_dim, curr_dim, (1, 1), bias=False),
nn.BatchNorm2d(curr_dim)
) for _ in range(nstack - 1)
])
self.cnvs_ = nn.ModuleList([
nn.Sequential(
nn.Conv2d(cnv_dim, curr_dim, (1, 1), bias=False),
nn.BatchNorm2d(curr_dim)
) for _ in range(nstack - 1)
])
self.tl_regrs = nn.ModuleList([
make_regr_layer(cnv_dim, curr_dim, 2) for _ in range(nstack)
])
self.br_regrs = nn.ModuleList([
make_regr_layer(cnv_dim, curr_dim, 2) for _ in range(nstack)
])
self.ct_regrs = nn.ModuleList([
make_regr_layer(cnv_dim, curr_dim, 2) for _ in range(nstack)
])
self.relu = nn.ReLU(inplace=True)
def _test(self, *xs, **kwargs):
image = xs[0]
inter = self.pre(image)
outs = []
layers = zip(
self.kps, self.cnvs,
self.tl_cnvs, self.br_cnvs,
self.ct_cnvs, self.tl_heats,
self.br_heats, self.ct_heats,
self.tl_tags, self.br_tags,
self.tl_regrs, self.br_regrs,
self.ct_regrs
)
for ind, layer in enumerate(layers):
kp_, cnv_ = layer[0:2]
tl_cnv_, br_cnv_ = layer[2:4]
ct_cnv_, tl_heat_ = layer[4:6]
br_heat_, ct_heat_ = layer[6:8]
tl_tag_, br_tag_ = layer[8:10]
tl_regr_, br_regr_ = layer[10:12]
ct_regr_ = layer[12]
kp = kp_(inter)
cnv = cnv_(kp)
if ind == self.nstack - 1:
tl_cnv = tl_cnv_(cnv)
br_cnv = br_cnv_(cnv)
ct_cnv = ct_cnv_(cnv)
tl_heat, br_heat, ct_heat = tl_heat_(tl_cnv), br_heat_(br_cnv), ct_heat_(ct_cnv)
tl_tag, br_tag = tl_tag_(tl_cnv), br_tag_(br_cnv)
tl_regr, br_regr, ct_regr = tl_regr_(tl_cnv), br_regr_(br_cnv), ct_regr_(ct_cnv)
outs += [tl_heat, br_heat, tl_tag, br_tag, tl_regr, br_regr,
ct_heat, ct_regr]
if ind < self.nstack - 1:
inter = self.inters_[ind](inter) + self.cnvs_[ind](cnv)
inter = self.relu(inter)
inter = self.inters[ind](inter)
return self._decode(*outs[-8:], **kwargs)
def forward(self, *xs, **kwargs):
if len(xs) > 1:
return self._train(*xs, **kwargs)
return self._test(*xs, **kwargs)
class pool(nn.Module):
def __init__(self, dim, pool1, pool2):
super(pool, self).__init__()
self.p1_conv1 = convolution(3, dim, 128)
self.p2_conv1 = convolution(3, dim, 128)
self.p_conv1 = nn.Conv2d(128, dim, (3, 3), padding=(1, 1), bias=False)
self.p_bn1 = nn.BatchNorm2d(dim)
self.conv1 = nn.Conv2d(dim, dim, (1, 1), bias=False)
self.bn1 = nn.BatchNorm2d(dim)
self.relu1 = nn.ReLU(inplace=True)
self.conv2 = convolution(3, dim, dim)
self.pool1 = pool1()
self.pool2 = pool2()
self.look_conv1 = convolution(3, dim, 128)
self.look_conv2 = convolution(3, dim, 128)
self.P1_look_conv = nn.Conv2d(128, 128, (3, 3), padding=(1, 1), bias=False)
self.P2_look_conv = nn.Conv2d(128, 128, (3, 3), padding=(1, 1), bias=False)
def forward(self, x):
# pool 1
look_conv1 = self.look_conv1(x)
p1_conv1 = self.p1_conv1(x)
look_right = self.pool2(look_conv1)
P1_look_conv = self.P1_look_conv(p1_conv1+look_right)
pool1 = self.pool1(P1_look_conv)
# pool 2
look_conv2 = self.look_conv2(x)
p2_conv1 = self.p2_conv1(x)
look_down = self.pool1(look_conv2)
P2_look_conv = self.P2_look_conv(p2_conv1+look_down)
pool2 = self.pool2(P2_look_conv)
# pool 1 + pool 2
p_conv1 = self.p_conv1(pool1 + pool2)
p_bn1 = self.p_bn1(p_conv1)
conv1 = self.conv1(x)
bn1 = self.bn1(conv1)
relu1 = self.relu1(p_bn1 + bn1)
conv2 = self.conv2(relu1)
return conv2
class pool_cross(nn.Module):
def __init__(self, dim, pool1, pool2, pool3, pool4):
super(pool_cross, self).__init__()
self.p1_conv1 = convolution(3, dim, 128)
self.p2_conv1 = convolution(3, dim, 128)
self.p_conv1 = nn.Conv2d(128, dim, (3, 3), padding=(1, 1), bias=False)
self.p_bn1 = nn.BatchNorm2d(dim)
self.conv1 = nn.Conv2d(dim, dim, (1, 1), bias=False)
self.bn1 = nn.BatchNorm2d(dim)
self.relu1 = nn.ReLU(inplace=True)
self.conv2 = convolution(3, dim, dim)
self.pool1 = pool1()
self.pool2 = pool2()
self.pool3 = pool3()
self.pool4 = pool4()
def forward(self, x):
# pool 1
p1_conv1 = self.p1_conv1(x)
pool1 = self.pool1(p1_conv1)
pool1 = self.pool3(pool1)
# pool 2
p2_conv1 = self.p2_conv1(x)
pool2 = self.pool2(p2_conv1)
pool2 = self.pool4(pool2)
# pool 1 + pool 2
p_conv1 = self.p_conv1(pool1 + pool2)
p_bn1 = self.p_bn1(p_conv1)
conv1 = self.conv1(x)
bn1 = self.bn1(conv1)
relu1 = self.relu1(p_bn1 + bn1)
conv2 = self.conv2(relu1)
return conv2
class tl_pool(pool):
def __init__(self, dim):
super(tl_pool, self).__init__(dim, TopPool, LeftPool)
class br_pool(pool):
def __init__(self, dim):
super(br_pool, self).__init__(dim, BottomPool, RightPool)
class center_pool(pool_cross):
def __init__(self, dim):
super(center_pool, self).__init__(dim, TopPool, LeftPool, BottomPool, RightPool)
def make_tl_layer(dim):
return tl_pool(dim)
def make_br_layer(dim):
return br_pool(dim)
def make_ct_layer(dim):
return center_pool(dim)
def make_pool_layer(dim):
return nn.Sequential()
def make_hg_layer(kernel, dim0, dim1, mod, layer=convolution, **kwargs):
layers = [layer(kernel, dim0, dim1, stride=2)]
layers += [layer(kernel, dim1, dim1) for _ in range(mod - 1)]
return nn.Sequential(*layers)
class CenterNet(kp):
def __init__(self):
n = 5
dims = [256, 256, 384, 384, 384, 512]
modules = [2, 2, 2, 2, 2, 4]
out_dim = 80
super(CenterNet, self).__init__(
n, 2, dims, modules, out_dim,
make_tl_layer=make_tl_layer,
make_br_layer=make_br_layer,
make_ct_layer=make_ct_layer,
make_pool_layer=make_pool_layer,
make_hg_layer=make_hg_layer,
kp_layer=residual, cnv_dim=256
)
torch.manual_seed(317)
# for model backward compatibility
# previously model was wrapped by DataParallel module
class DummyModule(nn.Module):
def __init__(self, model):
super(DummyModule, self).__init__()
self.module = model
def forward(self, *xs, **kwargs):
return self.module(*xs, **kwargs)
class NetworkFactory(object):
def __init__(self):
super(NetworkFactory, self).__init__()
self.model = CenterNet()
self.network = DataParallel(self.model, chunk_sizes=system_configs.chunk_sizes).cuda()
self.network.eval()
self.network.load_state_dict(torch.load(r'C:\Users\Administrator\Music\CenterNet\cache\nnet\CenterNet-104\CenterNet-104_480000.pkl'))
def test(self, xs, **kwargs):
return self.network(xs, **kwargs)
def parse_args():
parser = argparse.ArgumentParser(description="Test CenterNet")
parser.add_argument("cfg_file", help="config file", type=str)
parser.add_argument("--testiter", dest="testiter",
help="test at iteration i",
default=None, type=int)
parser.add_argument("--split", dest="split",
help="which split to use",
default="validation", type=str)
parser.add_argument("--suffix", dest="suffix", default=None, type=str)
parser.add_argument("--debug", action="store_true")
args = parser.parse_args()
return args
def _rescale_dets(detections, ratios, borders, sizes):
xs, ys = detections[..., 0:4:2], detections[..., 1:4:2]
xs /= ratios[:, 1][:, None, None]
ys /= ratios[:, 0][:, None, None]
xs -= borders[:, 2][:, None, None]
ys -= borders[:, 0][:, None, None]
tx_inds = xs[:, :, 0] <= -5
bx_inds = xs[:, :, 1] >= sizes[0, 1] + 5
ty_inds = ys[:, :, 0] <= -5
by_inds = ys[:, :, 1] >= sizes[0, 0] + 5
np.clip(xs, 0, sizes[:, 1][:, None, None], out=xs)
np.clip(ys, 0, sizes[:, 0][:, None, None], out=ys)
detections[:, tx_inds[0, :], 4] = -1
detections[:, bx_inds[0, :], 4] = -1
detections[:, ty_inds[0, :], 4] = -1
detections[:, by_inds[0, :], 4] = -1
def kp_decode(nnet, images, K, ae_threshold=0.5, kernel=3):
detections, center = nnet.test(images, ae_threshold=ae_threshold, K=K, kernel=kernel)
detections = detections.data.cpu().numpy()
center = center.data.cpu().numpy()
return detections, center
LeiBie = [ "person",
"bicycle",
"car",
"motorcycle",
"airplane",
"bus",
"train",
"truck",
"boat",
"traffic light",
"fire hydrant",
"stop sign",
"parking meter",
"bench",
"bird",
"cat",
"dog",
"horse",
"sheep",
"cow",
"elephant",
"bear",
"zebra",
"giraffe",
"backpack",
"umbrella",
"handbag",
"tie",
"suitcase",
"frisbee",
"skis",
"snowboard",
"sports ball",
"kite",
"baseball bat",
"baseball glove",
"skateboard",
"surfboard",
"tennis racket",
"bottle",
"wine glass",
"cup",
"fork",
"knife",
"spoon",
"bowl",
"banana",
"apple",
"sandwich",
"orange",
"broccoli",
"carrot",
"hot dog",
"pizza",
"donut",
"cake",
"chair",
"couch",
"potted plant",
"bed",
"dining table",
"toilet",
"tv",
"laptop",
"mouse",
"remote",
"keyboard",
"cell phone",
"microwave",
"oven",
"toaster",
"sink",
"refrigerator",
"book",
"clock",
"vase",
"scissors",
"teddy bear",
"hair drier",
"toothbrush"]
def kp_detection(nnet, image,decode_func=kp_decode):
K = 100
ae_threshold = 0.5
nms_kernel = 3
scales = [1]
weight_exp = 8
merge_bbox = False
categories = 80
nms_threshold = 0.5
max_per_image = 100
nms_algorithm = {
"nms": 0,
"linear_soft_nms": 1,
"exp_soft_nms": 2
}["exp_soft_nms"]
num_images = 1
top_bboxes = {}
for ind in tqdm(range(0, num_images), ncols=80, desc="locating kps"):
height, width = image.shape[0:2]
detections = []
center_points = []
for scale in scales:
new_height = int(height * scale)
new_width = int(width * scale)
new_center = np.array([new_height // 2, new_width // 2])
inp_height = new_height | 127
inp_width = new_width | 127
images = np.zeros((1, 3, inp_height, inp_width), dtype=np.float32)
ratios = np.zeros((1, 2), dtype=np.float32)
borders = np.zeros((1, 4), dtype=np.float32)
sizes = np.zeros((1, 2), dtype=np.float32)
out_height, out_width = (inp_height + 1) // 4, (inp_width + 1) // 4
height_ratio = out_height / inp_height
width_ratio = out_width / inp_width
resized_image = cv2.resize(image, (new_width, new_height))
resized_image, border, offset = crop_image(resized_image, new_center, [inp_height, inp_width])
resized_image = resized_image / 255
images[0] = resized_image.transpose((2, 0, 1))
borders[0] = border
sizes[0] = [int(height * scale), int(width * scale)]
ratios[0] = [height_ratio, width_ratio]
images = np.concatenate((images, images[:, :, :, ::-1]), axis=0)
images = torch.from_numpy(images)
dets, center = decode_func(nnet, images, K, ae_threshold=ae_threshold, kernel=nms_kernel)
dets = dets.reshape(2, -1, 8)
center = center.reshape(2, -1, 4)
dets[1, :, [0, 2]] = out_width - dets[1, :, [2, 0]]
center[1, :, [0]] = out_width - center[1, :, [0]]
dets = dets.reshape(1, -1, 8)
center = center.reshape(1, -1, 4)
_rescale_dets(dets, ratios, borders, sizes)
center[..., [0]] /= ratios[:, 1][:, None, None]
center[..., [1]] /= ratios[:, 0][:, None, None]
center[..., [0]] -= borders[:, 2][:, None, None]
center[..., [1]] -= borders[:, 0][:, None, None]
np.clip(center[..., [0]], 0, sizes[:, 1][:, None, None], out=center[..., [0]])
np.clip(center[..., [1]], 0, sizes[:, 0][:, None, None], out=center[..., [1]])
dets[:, :, 0:4] /= scale
center[:, :, 0:2] /= scale
if scale == 1:
center_points.append(center)
detections.append(dets)
detections = np.concatenate(detections, axis=1)
center_points = np.concatenate(center_points, axis=1)
classes = detections[..., -1]
classes = classes[0]
detections = detections[0]
center_points = center_points[0]
valid_ind = detections[:, 4] > -1
valid_detections = detections[valid_ind]
box_width = valid_detections[:, 2] - valid_detections[:, 0]
box_height = valid_detections[:, 3] - valid_detections[:, 1]
s_ind = (box_width * box_height <= 22500)
l_ind = (box_width * box_height > 22500)
s_detections = valid_detections[s_ind]
l_detections = valid_detections[l_ind]
s_left_x = (2 * s_detections[:, 0] + s_detections[:, 2]) / 3
s_right_x = (s_detections[:, 0] + 2 * s_detections[:, 2]) / 3
s_top_y = (2 * s_detections[:, 1] + s_detections[:, 3]) / 3
s_bottom_y = (s_detections[:, 1] + 2 * s_detections[:, 3]) / 3
s_temp_score = copy.copy(s_detections[:, 4])
s_detections[:, 4] = -1
center_x = center_points[:, 0][:, np.newaxis]
center_y = center_points[:, 1][:, np.newaxis]
s_left_x = s_left_x[np.newaxis, :]
s_right_x = s_right_x[np.newaxis, :]
s_top_y = s_top_y[np.newaxis, :]
s_bottom_y = s_bottom_y[np.newaxis, :]
ind_lx = (center_x - s_left_x) > 0
ind_rx = (center_x - s_right_x) < 0
ind_ty = (center_y - s_top_y) > 0
ind_by = (center_y - s_bottom_y) < 0
ind_cls = (center_points[:, 2][:, np.newaxis] - s_detections[:, -1][np.newaxis, :]) == 0
ind_s_new_score = np.max(((ind_lx + 0) & (ind_rx + 0) & (ind_ty + 0) & (ind_by + 0) & (ind_cls + 0)),
axis=0) == 1
index_s_new_score = np.argmax(
((ind_lx + 0) & (ind_rx + 0) & (ind_ty + 0) & (ind_by + 0) & (ind_cls + 0))[:, ind_s_new_score], axis=0)
s_detections[:, 4][ind_s_new_score] = (s_temp_score[ind_s_new_score] * 2 + center_points[
index_s_new_score, 3]) / 3
l_left_x = (3 * l_detections[:, 0] + 2 * l_detections[:, 2]) / 5
l_right_x = (2 * l_detections[:, 0] + 3 * l_detections[:, 2]) / 5
l_top_y = (3 * l_detections[:, 1] + 2 * l_detections[:, 3]) / 5
l_bottom_y = (2 * l_detections[:, 1] + 3 * l_detections[:, 3]) / 5
l_temp_score = copy.copy(l_detections[:, 4])
l_detections[:, 4] = -1
center_x = center_points[:, 0][:, np.newaxis]
center_y = center_points[:, 1][:, np.newaxis]
l_left_x = l_left_x[np.newaxis, :]
l_right_x = l_right_x[np.newaxis, :]
l_top_y = l_top_y[np.newaxis, :]
l_bottom_y = l_bottom_y[np.newaxis, :]
ind_lx = (center_x - l_left_x) > 0
ind_rx = (center_x - l_right_x) < 0
ind_ty = (center_y - l_top_y) > 0
ind_by = (center_y - l_bottom_y) < 0
ind_cls = (center_points[:, 2][:, np.newaxis] - l_detections[:, -1][np.newaxis, :]) == 0
ind_l_new_score = np.max(((ind_lx + 0) & (ind_rx + 0) & (ind_ty + 0) & (ind_by + 0) & (ind_cls + 0)),
axis=0) == 1
index_l_new_score = np.argmax(
((ind_lx + 0) & (ind_rx + 0) & (ind_ty + 0) & (ind_by + 0) & (ind_cls + 0))[:, ind_l_new_score], axis=0)
l_detections[:, 4][ind_l_new_score] = (l_temp_score[ind_l_new_score] * 2 + center_points[
index_l_new_score, 3]) / 3
detections = np.concatenate([l_detections, s_detections], axis=0)
detections = detections[np.argsort(-detections[:, 4])]
classes = detections[..., -1]
keep_inds = (detections[:, 4] > -1)
detections = detections[keep_inds]
classes = classes[keep_inds]#9个
image_id = '1'
top_bboxes[image_id] = {}
for j in range(categories):
keep_inds = (classes == j)
top_bboxes[image_id][j + 1] = detections[keep_inds][:, 0:7].astype(np.float32)
if merge_bbox:
soft_nms_merge(top_bboxes[image_id][j + 1], Nt=nms_threshold, method=nms_algorithm,
weight_exp=weight_exp)
else:
#执行的这个
soft_nms(top_bboxes[image_id][j + 1], Nt=nms_threshold, method=nms_algorithm)
top_bboxes[image_id][j + 1] = top_bboxes[image_id][j + 1][:, 0:5]
scores = np.hstack([
top_bboxes[image_id][j][:, -1]
for j in range(1, categories + 1)
])
#未执行，限制物体总数不超过100
if len(scores) > max_per_image:
kth = len(scores) - max_per_image
thresh = np.partition(scores, kth)[kth]
for j in range(1, categories + 1):
keep_inds = (top_bboxes[image_id][j][:, -1] >= thresh)
top_bboxes[image_id][j] = top_bboxes[image_id][j][keep_inds]
debug = True
if debug:
im = image[:, :, (2, 1, 0)]
fig, ax = plt.subplots(figsize=(12, 12))
fig = ax.imshow(im, aspect='equal')
plt.axis('off')
fig.axes.get_xaxis().set_visible(False)
fig.axes.get_yaxis().set_visible(False)
for j in range(1, categories + 1):
keep_inds = (top_bboxes[image_id][j][:, -1] >= 0.4)
cat_name = LeiBie[j - 1]
for bbox in top_bboxes[image_id][j][keep_inds]:
bbox = bbox[0:4].astype(np.int32)
xmin = bbox[0]
ymin = bbox[1]
xmax = bbox[2]
ymax = bbox[3]
# if (xmax - xmin) * (ymax - ymin) > 5184:
ax.add_patch(
plt.Rectangle((xmin, ymin), xmax - xmin, ymax - ymin, fill=False, edgecolor=colours[j - 1],
linewidth=4.0))
ax.text(xmin + 1, ymin - 3, '{:s}'.format(cat_name),
bbox=dict(facecolor=colours[j - 1], ec='black', lw=2, alpha=0.5),
fontsize=15, color='white', weight='bold')
plt.show()
plt.waitforbuttonpress(0)
# debug_file1 = "a.pdf"
# debug_file2 = "a.jpg"
# plt.savefig(debug_file1)
# plt.savefig(debug_file2)
# plt.close()
if __name__ == "__main__":
nnet = NetworkFactory()
image = cv2.imread(r'D:\xunleixiazai\train2014\COCO_train2014_000000001912.jpg')
kp_detection(nnet=nnet,image=image)运行：python demo.py CenterNet-104
789
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
