518789Adf
commented
3 months ago class MaskRcnnDataset(Dataset): def init(self, annotation_lines, input_shape=[600, 600], train=True): self.annotation_lines = annotation_lines self.length = len(annotation_lines) self.input_shape = input_shape self.train = train
def __len__(self):
return self.length
def __getitem__(self, index):
index = index % self.length
#---------------------------------------------------#
# 训练时进行数据的随机增强
# 验证时不进行数据的随机增强
#---------------------------------------------------#
image, y, masks = self.get_random_data(self.annotation_lines[index], self.input_shape[0:2], random=self.train)
image = np.transpose(F.to_tensor(np.array(image, dtype=np.float32)), (2, 0, 1))
box_data = np.zeros((len(y), 5))
if len(y) > 0:
box_data[:len(y)] = y
box = box_data[:, :4]
label = box_data[:, -1]
# Convert masks to the correct format
masks = [Image.fromarray(mask.astype(np.uint8)) for mask in masks]
masks = [F.to_tensor(mask) for mask in masks]
masks = torch.stack(masks, dim=0)
return image, box, label, masks
def rand(self, a=0, b=1):
return np.random.rand() * (b - a) + a
def get_random_data(self, annotation_line, input_shape, jitter=.3, hue=.1, sat=0.7, val=0.4, random=True):
line = annotation_line.split()
#------------------------------#
# 读取图像并转换成RGB图像
#------------------------------#
image = Image.open(line[0])
image = cvtColor(image)
#------------------------------#
# 获得图像的高宽与目标高宽
#------------------------------#
iw, ih = image.size
h, w = input_shape
#------------------------------#
# 获得预测框
#------------------------------#
box = np.array([np.array(list(map(int, box.split(',')))) for box in line[1:]])
#------------------------------#
# 加载掩码数据
#------------------------------#
masks = [Image.open(line[i + 1]).convert('1') for i in range(len(box))]
if not random:
scale = min(w / iw, h / ih)
nw = int(iw * scale)
nh = int(ih * scale)
dx = (w - nw) // 2
dy = (h - nh) // 2
# 将图像多余的部分加上灰条
image = image.resize((nw, nh), Image.BICUBIC)
new_image = Image.new('RGB', (w, h), (128, 128, 128))
new_image.paste(image, (dx, dy))
image_data = np.array(new_image, np.float32)
# 对真实框进行调整
if len(box) > 0:
np.random.shuffle(box)
box[:, [0, 2]] = box[:, [0, 2]] * nw / iw + dx
box[:, [1, 3]] = box[:, [1, 3]] * nh / ih + dy
# 对掩码进行调整
masks = [mask.resize((nw, nh), Image.NEAREST) for mask in masks]
masks = [np.pad(mask, ((0, h - nh), (0, w - nw)), mode='constant') for mask in masks]
box[:, 0:2][box[:, 0:2] < 0] = 0
box[:, 2][box[:, 2] > w] = w
box[:, 3][box[:, 3] > h] = h
box_w = box[:, 2] - box[:, 0]
box_h = box[:, 3] - box[:, 1]
box = box[np.logical_and(box_w > 1, box_h > 1)] # discard invalid box
masks = [mask for mask, b in zip(masks, box) if b[2] - b[0] > 1 and b[3] - b[1] > 1]
return image_data, box, masks
# 对图像进行缩放并且进行长和宽的扭曲
new_ar = iw / ih * self.rand(1 - jitter, 1 + jitter) / self.rand(1 - jitter, 1 + jitter)
scale = self.rand(.25, 2)
if new_ar < 1:
nh = int(scale * h)
nw = int(nh * new_ar)
else:
nw = int(scale * w)
nh = int(nw / new_ar)
image = image.resize((nw, nh), Image.BICUBIC)
# 将图像多余的部分加上灰条
dx = int(self.rand(0, w - nw))
dy = int(self.rand(0, h - nh))
new_image = Image.new('RGB', (w, h), (128, 128, 128))
new_image.paste(image, (dx, dy))
image = new_image
# 翻转图像
flip = self.rand() < .5
if flip:
image = image.transpose(Image.FLIP_LEFT_RIGHT)
masks = [mask.transpose(Image.FLIP_LEFT_RIGHT) for mask in masks]
image_data = np.array(image, np.uint8)
# 对图像进行色域变换
r = np.random.uniform(-1, 1, 3) * [hue, sat, val] + 1
# 将图像转到HSV上
hue, sat, val = cv2.split(cv2.cvtColor(image_data, cv2.COLOR_RGB2HSV))
dtype = image_data.dtype
# 应用变换
x = np.arange(0, 256, dtype=r.dtype)
lut_hue = ((x * r[0]) % 180).astype(dtype)
lut_sat = np.clip(x * r[1], 0, 255).astype(dtype)
lut_val = np.clip(x * r[2], 0, 255).astype(dtype)
image_data = cv2.merge((cv2.LUT(hue, lut_hue), cv2.LUT(sat, lut_sat), cv2.LUT(val, lut_val)))
image_data = cv2.cvtColor(image_data, cv2.COLOR_HSV2RGB)
# 对真实框进行调整
if len(box) > 0:
np.random.shuffle(box)
box[:, [0, 2]] = box[:, [0, 2]] * nw / iw + dx
box[:, [1, 3]] = box[:, [1, 3]] * nh / ih + dy
if flip:
box[:, [0, 2]] = w - box[:, [2, 0]]
box[:, 0:2][box[:, 0:2] < 0] = 0
box[:, 2][box[:, 2] > w] = w
box[:, 3][box[:, 3] > h] = h
box_w = box[:, 2] - box[:, 0]
box_h = box[:, 3] - box[:, 1]
box = box[np.logical_and(box_w > 1, box_h > 1)]
# 对掩码进行调整
masks = [mask.resize((nw, nh), Image.NEAREST) for mask in masks]
masks = [np.pad(mask, ((0, h - nh), (0, w - nw)), mode='constant') for mask in masks]
masks = [mask.transpose(Image.FLIP_LEFT_RIGHT) if flip else mask for mask in masks]
masks = [mask for mask, b in zip(masks, box) if b[2] - b[0] > 1 and b[3] - b[1] > 1]
return image_data, box, masks
有代码，以下是maskrcnn在这个基代码的实现： class MaskRCNNHead(nn.Module): def init(self, n_class, roi_size, spatial_scale, num_convs=4, conv_dim=256, mask_out_dim=28): super(MaskRCNNHead, self).init()
def roi_pooling(features, roi, roi_index, roi_size, spatial_scale):
假设roi_pooling是一个已经定义好的函数，用于从特征图中裁剪出对应ROI的部分
还有就是，再建立一个文件maskrcnn.py： import torch.nn as nn
from nets.classifier import Resnet50RoIHead, VGG16RoIHead, MaskRCNNHead from nets.resnet50 import resnet50 from nets.rpn import RegionProposalNetwork from nets.vgg16 import decom_vgg16
class MaskRCNN(): def init(self, num_classes,
mode = "training", feat_stride = 16, anchor_scales = [8, 16, 32], ratios = [0.5, 1, 2], backbone = 'vgg', pretrained = False): super(MaskRCNN, self).init() self.feat_stride = feat_stride
---------------------------------
即可完成maskrcnn的训练功能！