目标检测是计算机视觉中一项非常流行的任务，在给定图像的情况下，您可以预测图像中存在的对象周围的（通常是矩形）框，并识别对象的类型。您的图像中可能有多个对象，并且有各种最先进的技术和架构来解决这个问题，例如Faster-RCNN和YOLO v3+。
本文讨论的是图像中只有一个感兴趣的对象的情况。这里的重点更多的是如何读取图像及其边界框，调整大小和正确地执行数据增强，而不是模型本身。我们的目标是很好地掌握对象检测背后的基本思想，您可以对其进行扩展，以便更好地理解更复杂的技术。
文末包含了我在本文中使用的所有代码。
1.问题陈述
给定一个由路标组成的图像，预测路标周围的边界框，并识别路标的类型。
这些路标可以分为四个不同的类别:
- Traffic Light
- Stop
- Speed Limit
- Crosswalk
这被称为多任务学习问题，因为它涉及执行两个任务:
- 1)回归找到边界框坐标
- 2)分类识别路标类型
2.数据集
我使用了Kaggle的路标检测数据集:
https://www.kaggle.com/andrewmvd/road-sign-detection
它由
877
张图片组成。这是一个相当不平衡的数据集，大多数图像属于
Speed Limit
类，但由于我们更关注边界框预测，我们可以忽略数据不平衡。
3.加载数据
每个图像的注释存储在单独的XML文件中。我按照以下步骤创建训练数据帧:
-
遍历训练目录以获得所有
.xml
文件的列表。 -
使用
xml.etree.ElementTree
解析
.xml
文件。 -
为每个图像创建一个包含
filepath、width 、height 、边界框坐标(xmin、xmax、ymin、ymax)和class
的字典，并将字典附加到列表中。 - 使用图像统计的字典列表创建一个pandas dataframe。
def filelist(root, file_type):
"""Returns a fully-qualified list of filenames under root directory"""
return [os.path.join(directory_path, f) for directory_path, directory_name,
files in os.walk(root) for f in files if f.endswith(file_type)]
def generate_train_df (anno_path):
annotations = filelist(anno_path, '.xml')
anno_list = []
for anno_path in annotations:
root = ET.parse(anno_path).getroot()
anno = {}
anno['filename'] = Path(str(images_path) + '/'+ root.find("./filename").text)
anno['width'] = root.find("./size/width").text
anno['height'] = root.find("./size/height").text
anno['class'] = root.find("./object/name").text
anno['xmin'] = int(root.find("./object/bndbox/xmin").text)
anno['ymin'] = int(root.find("./object/bndbox/ymin").text)
anno['xmax'] = int(root.find("./object/bndbox/xmax").text)
anno['ymax'] = int(root.find("./object/bndbox/ymax").text)
anno_list.append(anno)
return pd.DataFrame(anno_list)
-
标签编码
class
列
#label encode target
class_dict = {'speedlimit': 0, 'stop': 1, 'crosswalk': 2, 'trafficlight': 3}
df_train['class'] = df_train['class'].apply(lambda x: class_dict[x])
4.调整图像和包围框的大小
由于训练计算机视觉模型需要图像具有相同的大小，所以我们需要调整图像及其相应的边界框的大小。调整图像的大小很简单，但调整边界框的大小有点棘手，因为每个框都是相对于图像及其尺寸的。
以下是调整边界框大小的方法:
- 根据目标图像的大小调整图像，并根据边界框获得掩码
- 将掩码的大小调整为所需的尺寸。这个掩码的背景值为0，边框覆盖的区域值为1。
- 从调整后的掩码中提取边界框坐标。
def create_mask(bb, x):
"""Creates a mask for the bounding box of same shape as image"""
rows,cols,*_ = x.shape
Y = np.zeros((rows, cols))
bb = bb.astype(np.int)
Y[bb[0]:bb[2], bb[1]:bb[3]] = 1.
return Y
def mask_to_bb(Y):
"""Convert mask Y to a bounding box, assumes 0 as background nonzero object"""
cols, rows = np.nonzero(Y)
if len(cols)==0:
return np.zeros(4, dtype=np.float32)
top_row = np.min(rows)
left_col = np.min(cols)
bottom_row = np.max(rows)
right_col = np.max(cols)
return np.array([left_col, top_row, right_col, bottom_row], dtype=np.float32)
def create_bb_array(x):
"""Generates bounding box array from a train_df row"""
return np.array([x[5],x[4],x[7],x[6]])
def resize_image_bb(read_path,write_path,bb,sz):
"""Resize an image and its bounding box and write image to new path"""
im = read_image(read_path)
im_resized = cv2.resize(im, (int(1.49*sz), sz))
Y_resized = cv2.resize(create_mask(bb, im), (int(1.49*sz), sz))
new_path = str(write_path/read_path.parts[-1])
cv2.imwrite(new_path, cv2.cvtColor(im_resized, cv2.COLOR_RGB2BGR))
return new_path, mask_to_bb(Y_resized)
#Populating Training DF with new paths and bounding boxes
new_paths = []
new_bbs = []
train_path_resized = Path('./road_signs/images_resized')
for index, row in df_train.iterrows():
new_path,new_bb = resize_image_bb(row['filename'], train_path_resized, create_bb_array(row.values),300)
new_paths.append(new_path)
new_bbs.append(new_bb)
df_train['new_path'] = new_paths
df_train['new_bb'] = new_bbs
5.数据增强
数据增强是一种更好地泛化我们模型的技术，它通过使用现有图像的不同变体来创建新的训练图像。我们目前的训练集中只有800张图像，所以数据增强对于确保我们的模型不会过度拟合非常重要。
对于这个问题，我使用了翻转，旋转，中心裁剪和随机裁剪。
这里唯一要记住的是确保边界框也以与图像相同的方式转换。要做到这一点，我们遵循与调整大小相同的方法——将边界框转换为一个掩码，对掩码应用与原始图像相同的转换，并提取边界框坐标。
# modified from fast.ai
def crop(im, r, c, target_r, target_c):
return im[r:r+target_r, c:c+target_c]
# random crop to the original size
def random_crop(x, r_pix=8):
""" Returns a random crop"""
r, c,*_ = x.shape
c_pix = round(r_pix*c/r)
rand_r = random.uniform(0, 1)
rand_c = random.uniform(0, 1)
start_r = np.floor(2*rand_r*r_pix).astype(int)
start_c = np.floor(2*rand_c*c_pix).astype(int)
return crop(x, start_r, start_c, r-2*r_pix, c-2*c_pix)
def center_crop(x, r_pix=8):
r, c,*_ = x.shape
c_pix = round(r_pix*c/r)
return crop(x, r_pix, c_pix, r-2*r_pix, c-2*c_pix)
def rotate_cv(im, deg, y=False, mode=cv2.BORDER_REFLECT, interpolation=cv2.INTER_AREA):
""" Rotates an image by deg degrees"""
r,c,*_ = im.shape
M = cv2.getRotationMatrix2D((c/2,r/2),deg,1)
if y:
return cv2.warpAffine(im, M,(c,r), borderMode=cv2.BORDER_CONSTANT)
return cv2.warpAffine(im,M,(c,r), borderMode=mode, flags=cv2.WARP_FILL_OUTLIERS+interpolation)
def random_cropXY(x, Y, r_pix=8):
""" Returns a random crop"""
r, c,*_ = x.shape
c_pix = round(r_pix*c/r)
rand_r = random.uniform(0, 1)
rand_c = random.uniform(0, 1)
start_r = np.floor(2*rand_r*r_pix).astype(int)
start_c = np.floor(2*rand_c*c_pix).astype(int)
xx = crop(x, start_r, start_c, r-2*r_pix, c-2*c_pix)
YY = crop(Y, start_r, start_c, r-2*r_pix, c-2*c_pix)
return xx, YY
def transformsXY(path, bb, transforms):
x = cv2.imread(str(path)).astype(np.float32)
x = cv2.cvtColor(x, cv2.COLOR_BGR2RGB)/255
Y = create_mask(bb, x)
if transforms:
rdeg = (np.random.random()-.50)*20
x = rotate_cv(x, rdeg)
Y = rotate_cv(Y, rdeg, y=True)
if np.random.random() > 0.5:
x = np.fliplr(x).copy()
Y = np.fliplr(Y).copy()
x, Y = random_cropXY(x, Y)
else:
x, Y = center_crop(x), center_crop(Y)
return x, mask_to_bb(Y)
def create_corner_rect(bb, color='red'):
bb = np.array(bb, dtype=np.float32)
return plt.Rectangle((bb[1], bb[0]), bb[3]-bb[1], bb[2]-bb[0], color=color,
fill=False, lw=3)
def show_corner_bb(im, bb):
plt.imshow(im)
plt.gca().add_patch(create_corner_rect(bb))
#original
im = cv2.imread(str(df_train.values[68][8]))
im = cv2.cvtColor(im, cv2.COLOR_BGR2RGB)
show_corner_bb(im, df_train.values[68][9])
# after transformation
im, bb = transformsXY(str(df_train.values[68][8]),df_train.values[68][9],True )
show_corner_bb(im, bb)
6.PyTorch Dataset
现在我们已经完成了数据增强，我们可以进行训练验证分离并创建PyTorch dataset。我们使用ImageNet统计数据对图像进行规范化，因为我们使用预训练的ResNet模型，并在训练时在数据集中应用数据增强。
X_train, X_val, y_train, y_val = train_test_split(X, Y, test_size=0.2, random_state=42)
def normalize(im):
"""Normalizes images with Imagenet stats."""
imagenet_stats = np.array([[0.485, 0.456, 0.406], [0.229, 0.224, 0.225]])
return (im - imagenet_stats[0])/imagenet_stats[1]
class RoadDataset(Dataset):
def __init__(self, paths, bb, y, transforms=False):
self.transforms = transforms
self.paths = paths.values
self.bb = bb.values
self.y = y.values
def __len__(self):
return len(self.paths)
def __getitem__(self, idx):
path = self.paths[idx]
y_class = self.y[idx]
x, y_bb = transformsXY(path, self.bb[idx], self.transforms)
x = normalize(x)
x = np.rollaxis(x, 2)
return x, y_class, y_bb
train_ds = RoadDataset(X_train['new_path'],X_train['new_bb'] ,y_train, transforms=True)
valid_ds = RoadDataset(X_val['new_path'],X_val['new_bb'],y_val)
batch_size = 32
train_dl = DataLoader(train_ds, batch_size=batch_size, shuffle=True)
valid_dl = DataLoader(valid_ds, batch_size=batch_size)
7.PyTorch模型
对于这个模型，我使用了一个非常简单的预训练resNet-34模型。因为我们在这里有两个任务要完成，所以最后还有两个层——边界框回归器和图像分类器。
class BB_model(nn.Module):
def __init__(self):
super(BB_model, self).__init__()
resnet = models.resnet34(pretrained=True)
layers = list(resnet.children())[:8]
self.features1 = nn.Sequential(*layers[:6])
self.features2 = nn.Sequential(*layers[6:])
self.classifier = nn.Sequential(nn.BatchNorm1d(512), nn.Linear(512, 4))
self.bb = nn.Sequential(nn.BatchNorm1d(512), nn.Linear(512, 4))
def forward(self, x):
x = self.features1(x)
x = self.features2(x)
x = F.relu(x)
x = nn.AdaptiveAvgPool2d((1,1))(x)
x = x.view(x.shape[0], -1)
return self.classifier(x), self.bb(x)
8.训练
对于损失，我们需要同时考虑分类损失和边界框回归损失，所以我们使用交叉熵和L1损失(真实值与预测坐标之间的所有绝对差值之和)的组合。我将L1损失缩小了1000倍，因为分类损失和回归损失的范围都差不多。除此之外，这是一个标准的PyTorch训练循环(使用GPU):
def update_optimizer(optimizer, lr):
for i, param_group in enumerate(optimizer.param_groups):
param_group["lr"] = lr
def train_epocs(model, optimizer, train_dl, val_dl, epochs=10,C=1000):
idx = 0
for i in range(epochs):
model.train()
total = 0
sum_loss = 0
for x, y_class, y_bb in train_dl:
batch = y_class.shape[0]
x = x.cuda().float()
y_class = y_class.cuda()
y_bb = y_bb.cuda().float()
out_class, out_bb = model(x)
loss_class = F.cross_entropy(out_class, y_class, reduction="sum")
loss_bb = F.l1_loss(out_bb, y_bb, reduction="none").sum(1)
loss_bb = loss_bb.sum()
loss = loss_class + loss_bb/C
optimizer.zero_grad()
loss.backward()
optimizer.step()
idx += 1
total += batch
sum_loss += loss.item()
train_loss = sum_loss/total
val_loss, val_acc = val_metrics(model, valid_dl, C)
print("train_loss %.3f val_loss %.3f val_acc %.3f" % (train_loss, val_loss, val_acc))
return sum_loss/total
def val_metrics(model, valid_dl, C=1000):
model.eval()
total = 0
sum_loss = 0
correct = 0
for x, y_class, y_bb in valid_dl:
batch = y_class.shape[0]
x = x.cuda().float()
y_class = y_class.cuda()
y_bb = y_bb.cuda().float()
out_class, out_bb = model(x)
loss_class = F.cross_entropy(out_class, y_class, reduction="sum")
loss_bb = F.l1_loss(out_bb, y_bb, reduction="none").sum(1)
loss_bb = loss_bb.sum()
loss = loss_class + loss_bb/C
_, pred = torch.max(out_class, 1)
correct += pred.eq(y_class).sum().item()
sum_loss += loss.item()
total += batch
return sum_loss/total, correct/total
model = BB_model().cuda()
parameters = filter(lambda p: p.requires_grad, model.parameters())
optimizer = torch.optim.Adam(parameters, lr=0.006)
train_epocs(model, optimizer, train_dl, valid_dl, epochs=15)
# train_loss 2.476 val_loss 51187.074 val_acc 0.119
# train_loss 1.661 val_loss 998.014 val_acc 0.091
# train_loss 1.657 val_loss 54.162 val_acc 0.261
9.对测试图像的预测
现在我们已经完成了训练，我们可以选择一个随机的图像并测试我们的模型。即使我们有相当少的训练图像，我们最终得到一个对我们的测试图像相当不错的预测。
用你的手机拍一张真实的照片并测试一下这个模型将是一项有趣的练习。另一个有趣的实验是不进行任何数据扩充，而是训练模型并比较两个模型。
# resizing test image
im = read_image('./road_signs/images_resized/road789.png')
im = cv2.resize(im, (int(1.49*300), 300))
cv2.imwrite('./road_signs/road_signs_test/road789.jpg', cv2.cvtColor(im, cv2.COLOR_RGB2BGR))
# test Dataset
test_ds = RoadDataset(pd.DataFrame([{'path':'./road_signs/road_signs_test/road789.jpg'}])['path'],pd.DataFrame([{'bb':np.array([0,0,0,0])}])['bb'],pd.DataFrame([{'y':[0]}])['y'])
x, y_class, y_bb = test_ds[0]
xx = torch.FloatTensor(x[None,])
print(xx.shape)
# torch.Size([1, 3, 284, 423])
# prediction
out_class, out_bb = model(xx.cuda())
print(out_class, out_bb)
# predicted bounding box
bb_hat = out_bb.detach().cpu().numpy()
bb_hat = bb_hat.astype(int)
show_corner_bb(im, bb_hat[0])
总结
现在我们已经介绍了对象检测的基本原理并从头开始实现它，您可以将这些思想扩展到多对象情况，并尝试更复杂的模型，如
RCNN和YOLO
!另外，可以使用这个名为
albumentations
的超级酷的库轻松地执行数据增强。
链接：https://pan.baidu.com/s/1euTGnDM8s4NED_5zMlTeCQ
提取码：123a
参考目录
https://towardsdatascience.com/bounding-box-prediction-from-scratch-using-pytorch-a8525da51ddc
版权声明：本文为weixin_43229348原创文章，遵循 CC 4.0 BY-SA 版权协议，转载请附上原文出处链接和本声明。