【AI达人特训营第三期】全球小麦麦穗检测

bboxes[:, 0] = bboxes[:, 0] + bboxes[:, 2] / 2.0  # 将真实框位置由左上转为中心点的xywh格式bboxes[:, 1] = bboxes[:, 1] + bboxes[:, 3] / 2.0

# 解压数据集!unzip -q -d data data/data198878/global-wheat-detection.zip

# 导入模块import numpy as npimport pandas as pdimport paddleimport osimport cv2from PIL import Image, ImageDraw, ImageEnhancefrom paddle.vision import transforms as Timport matplotlib.pyplot as pltfrom tqdm.notebook import tqdm # 进度条模块

ROOT_PATH = 'data' def get_path(*args,fp_postfix=None):
    # 获取文件路径
    '''
    Params: fp_postfix 文件后缀 如：jpg、png、gif等
    Params: *args 获取路径参数
    example:
        f = get_path('a', 'b', 'c', 'd')
        print(f) # data/a/b/c/d
        f = get_path('a', 'b', 'c', 'd',fp_postfix='jpg')
        print(f) # data/a/b/c/d.jpg
    '''
    if fp_postfix:
        obj_path = os.path.join(ROOT_PATH, *args) + '.' + fp_postfix        return obj_path

    obj_path = os.path.join(ROOT_PATH, *args)    return obj_path

# 查看数据集信息img_nums = len(os.listdir(get_path('train')))print(f'train中的图片张数:{img_nums}') # 查看train中的图片张数labels =  pd.read_csv(get_path('train.csv')) # 读取train.csv中的数据print("\n训练集有效图片数目：{}".format(labels["image_id"].nunique()))print("训练的真实框gt_box数目：",labels.shape[0])
labels.head()

def group_boxes(group):
    # 将image_id相同的图片
    boundaries = group["bbox"].str.split(",", expand = True)
    boundaries[0] = boundaries[0].str.slice(start = 1) # 去掉bbox字符串的'['
    boundaries[3] = boundaries[3].str.slice(stop = -1) # 去掉bbox字符串的']'
    
    return boundaries.values.astype(float)

labels = labels.groupby("image_id").apply(group_boxes)

print('单张图片真实框形状：', labels["ffbf75e5b"].shape)
labels["ffbf75e5b"]

# 切分数据集train_image_ids = np.unique(labels.index.values)[0:3358]
val_image_ids = np.unique(labels.index.values)[3358:3373]

# 加载训图片,将尺寸1024,1024 转成 256,256,以便更快训练def load_image(image_id):
    img_path = get_path('train', image_id, fp_postfix='jpg')
    img = Image.open(img_path)
    img = img.resize((256, 256))    
    return np.asarray(img)

# 获取图像和真实框并存储为字典def get_data(image_ids):
    data_pixels = {} # 图像内容
    data_labels = {} # 真实框坐标

    for image_id in tqdm(image_ids):
        data_pixels[image_id] = load_image(image_id)
        data_labels[image_id] = labels[image_id].copy() / 4
    return data_pixels, data_labels

# 训练集数据train_pixels, train_labels = get_data(train_image_ids)

# 验证集数据val_pixels, val_labels = get_data(val_image_ids)

def draw_bboxes(image_id, bboxes, source = "train"):  
    img_path = get_path(source, image_id, fp_postfix="jpg")
    image = Image.open(img_path) # 读取图片
    image = image.resize((256,256)) # 调整尺寸256*256

    # image = transform()(image) # 测试图像增广#################
    
    draw = ImageDraw.Draw(image) # 实例化图片
            
    for bbox in bboxes: # 遍历bboxes
        draw_bbox(draw, bbox) # 画出bbox
    
    return np.asarray(image)def draw_bbox(draw, bbox): # 画方框函数
    x, y, width, height = bbox
    draw.rectangle([x, y, x + width, y + height], width = 2, outline='red')

def show_images(image_ids, bboxes, source = 'train'):
    # 多个图像多次调用此函数。
    pixels = []    
    for image_id in image_ids:
        pixels.append(
            draw_bboxes(image_id, bboxes[image_id], source)
        )
    
    num_of_images = len(image_ids)
    fig, axes = plt.subplots(1, num_of_images, figsize = (5 * num_of_images, 5 * num_of_images))    
    for i, image_pixels in enumerate(pixels):
        axes[i].imshow(image_pixels)

show_images(train_image_ids[0:2], train_labels)

tiny_bboxes = []for i, image_id in enumerate(train_image_ids):    for label in train_labels[image_id]:        if (label[2] * label[3]) <= 10 and label[2] * label[3] != 0:
            tiny_bboxes.append((image_id, i))            
print(str(len(tiny_bboxes)) + " 个微小边框")# print(tiny_bboxes)

huge_bboxes = []for i, image_id in enumerate(train_image_ids):    for label in train_labels[image_id]:        if label[2] * label[3] > 8000:
            huge_bboxes.append((image_id, i))            
print(str(len(huge_bboxes)) + " 个巨大边框")# print(huge_bboxes)

# 展示部分无麦穗的边框show_images(train_image_ids[19:21], train_labels)

# 抽取边框好的训练数据，重新生成训练集，不改变原有数据def clean_labels(train_image_ids, train_labels):
    good_labels = {}    
    for i, image_id in enumerate(train_image_ids):
        good_labels[image_id] = []        
        for j, label in enumerate(train_labels[image_id]):            # remove huge bbox
            if label[2] * label[3] > 8000 and i not in [1079, 1371, 2020]:                continue

            # remove tiny bbox
            elif label[2] < 5 or label[3] < 5:                continue
                
            else:
                good_labels[image_id].append(
                    train_labels[image_id][j]
                )                
    return good_labels

train_labels = clean_labels(train_image_ids, train_labels)

# 数据加载全整体代码# 导入模块import osimport numpy as npimport pandas as pdfrom tqdm.notebook import tqdm # 进度条模块import cv2from PIL import Image, ImageDraw, ImageEnhanceimport matplotlib.pyplot as pltimport paddlefrom paddle.vision import transforms as T# 数据集目录ROOT_PATH = 'data' # 输入图片尺寸IMG_SIZE = 416def get_path(*args,fp_postfix=None):
    # 获取文件路径
    '''
    Params: fp_postfix 文件后缀 如：jpg、png、gif等
    Params: *args 获取路径参数
    example:
        f = get_path('a', 'b', 'c', 'd')
        print(f) # data/a/b/c/d
        f = get_path('a', 'b', 'c', 'd',fp_postfix='jpg')
        print(f) # data/a/b/c/d.jpg
    '''
    if fp_postfix:
        obj_path = os.path.join(ROOT_PATH, *args) + '.' + fp_postfix        return obj_path

    obj_path = os.path.join(ROOT_PATH, *args)    return obj_path# bbox分组，转numpydef group_boxes(group):
    # 将image_id相同的图片
    boundaries = group["bbox"].str.split(",", expand = True)
    boundaries[0] = boundaries[0].str.slice(start = 1) # 去掉bbox字符串的'['
    boundaries[3] = boundaries[3].str.slice(stop = -1) # 去掉bbox字符串的']'
    
    return boundaries.values.astype(float)# 加载训图片,将尺寸1024,1024 转成 IMG_SIZE*IMG_SIZE,以便更快训练def load_image(image_id):
    img_path = get_path('train', image_id, fp_postfix='jpg')
    img = Image.open(img_path)
    img = img.resize((IMG_SIZE, IMG_SIZE))    
    return np.asarray(img)# 切分数据集def train_seq_val(ration_size):
    train_image_ids = np.unique(labels.index.values)[0:ration_size]
    val_image_ids = np.unique(labels.index.values)[ration_size:3373]    return train_image_ids, val_image_ids# 获取图像和真实框并存储为字典def get_data(image_ids):
    data_pixels = {} # 图像内容
    data_labels = {} # 真实框坐标

    for image_id in tqdm(image_ids):
        data_pixels[image_id] = load_image(image_id)
        data_labels[image_id] = labels[image_id].copy() / (1024 / IMG_SIZE)  # 缩放标签
    return data_pixels, data_labels# 抽取边框好的训练数据，重新生成训练集，不改变原有数据def clean_labels(train_image_ids, train_labels):
    good_labels = {}    
    for i, image_id in enumerate(train_image_ids):
        good_labels[image_id] = []        
        for j, label in enumerate(train_labels[image_id]):            # remove huge bbox
            if label[2] * label[3] > 8000 and i not in [1079, 1371, 2020]:                continue

            # remove tiny bbox
            elif label[2] < 5 or label[3] < 5:                continue
                
            else:
                good_labels[image_id].append(
                    train_labels[image_id][j])                
    return good_labels


labels = pd.read_csv(get_path('train.csv')) # 读取train.csv中的数labels = labels.groupby("image_id").apply(group_boxes) # 格式转换,真实框分组# 切分数据集train_image_ids, val_image_ids = train_seq_val(ration_size=3358)# 训练集数据train_pixels, train_labels = get_data(train_image_ids)# 验证集数据val_pixels, val_labels = get_data(val_image_ids)# 进一步优化真实框数据train_labels = clean_labels(train_image_ids, train_labels)

def get_bbox(self, gt_bbox):
        # 对于一般的检测任务来说，一张图片上往往会有多个目标物体
        # 设置参数MAX_NUM = 55， 即一张图片最多取55个真实框；如果真实
        # 框的数目少于55个，则将不足部分的gt_bbox的各项数值全设置为0
        MAX_NUM = 55
        gt_bbox2 = np.zeros((MAX_NUM, 4))
        gt_class2 = np.zeros((MAX_NUM,))    # 1分类   
        gt_bbox = np.array(gt_bbox)        for i in range(len(gt_bbox)):            if i >= MAX_NUM:                break
            gt_bbox2[i, :] = gt_bbox[i, :]
            gt_class2[i] = 0 # 一分类标签为0
                   
        return gt_bbox2, gt_class2

class MyDataset(paddle.io.Dataset):
    def __init__(self, image_ids, image_pixels, img_size, labels = None, mode = None):
        super(MyDataset, self).__init__()
        self.image_ids = image_ids
        self.image_pixels = image_pixels
        self.img_size = img_size
        self.labels = labels
        self.transform = None
        self.mode = mode        if self.mode == "train":
            self.transform =T.Compose([                      
                T.BrightnessTransform(0.4), # 亮度调节
                T.ContrastTransform(0.4),   # 对比度调节
                T.HueTransform(0.4),        # 色调
                # T.RandomErasing(),        # 随机擦除
                T.Normalize(mean=[123.675, 116.28, 103.53], std=[58.395, 57.12, 57.375],data_format='HWC'), # 标准化
                T.Transpose()             # 数据格式转换,Transpose默认参数(2,0,1)
            ])        if self.mode == "val":
            self.transform =T.Compose([
                T.Normalize(mean=[123.675, 116.28, 103.53], std=[58.395, 57.12, 57.375],data_format='HWC'),
                T.Transpose()
            ])     
    
    def get_bbox(self, gt_bbox):
        # 对于一般的检测任务来说，一张图片上往往会有多个目标物体
        # 设置参数MAX_NUM = 55， 即一张图片最多取55个真实框；如果真实
        # 框的数目少于55个，则将不足部分的gt_bbox的各项数值全设置为0
        MAX_NUM = 55
        gt_bbox2 = np.zeros((MAX_NUM, 4))
        gt_class2 = np.zeros((MAX_NUM,))    # 1分类   
        gt_bbox = np.array(gt_bbox)        for i in range(len(gt_bbox)):            if i >= MAX_NUM:                break
            gt_bbox2[i, :] = gt_bbox[i, :]
            gt_class2[i] = 0
                   
        return gt_bbox2, gt_class2    def __getitem__(self, index):
        image_id = self.image_ids[index]     

        X = self.image_pixels[image_id]

        w = X.shape[0]
        h = X.shape[1]
        
        bboxes, gt_labels = self.get_bbox(self.labels[image_id])
        box_idx = np.arange(bboxes.shape[0]) # 随机打乱真实框排列顺序
        np.random.shuffle(box_idx)
        gt_labels = gt_labels[box_idx]

        bboxes = bboxes[box_idx] # 真实框位置是左上点的xywh格式
        bboxes[:, 0] = bboxes[:, 0] + bboxes[:, 2] / 2.0  # 将真实框位置由左上转为中心点的xywh格式
        bboxes[:, 1] = bboxes[:, 1] + bboxes[:, 3] / 2.0
        
        y = bboxes        if self.transform:
            X = self.transform(X)        return X.astype('float32'), np.array(y, dtype = 'float32') / self.img_size, \
            np.array(gt_labels, dtype = 'int32'), np.array([h, w],dtype='int32')    def __len__(self):
        return len(self.image_ids)

train_dataset = MyDataset(train_image_ids,
                        train_pixels,
                        IMG_SIZE,
                        train_labels,
                        mode='train')

val_dataset = MyDataset(val_image_ids,
                        val_pixels,
                        IMG_SIZE,
                        val_labels,
                        mode='val')# 查看数据集形状d = paddle.io.DataLoader(train_dataset, batch_size=4, shuffle=True, num_workers=0)print(next(d())[0].shape, next(d())[1].shape)

    loss_obj = paddle.nn.fucntional.binary_cross_entropy_with_logits(pred_classification, label_classification)

import timeimport osimport paddlefrom work.model import YOLOv3from work.utils import new_anchors


IMG_SIZE = 416ANCHORS = new_anchors(IMG_SIZE)
ANCHOR_MASKS = [[6, 7, 8], [3, 4, 5], [0, 1, 2]]
IGNORE_THRESH = 0.7# 类别数NUM_CLASSES = 1# 训练轮数MAX_EPOCH = 10def get_lr(base_lr = 0.000125, lr_decay = 0.1):
    bd = [10000, 20000]
    lr = [base_lr, base_lr * lr_decay, base_lr * lr_decay * lr_decay]
    learning_rate = paddle.optimizer.lr.PiecewiseDecay(boundaries=bd, values=lr)    return learning_rateif __name__ == '__main__':    # 设置gpu
    paddle.device.set_device("gpu")    # 实例化数据集
    train_dataset = MyDataset(train_image_ids,
                        train_pixels,
                        IMG_SIZE,
                        train_labels,
                        mode='train')
    val_dataset = MyDataset(val_image_ids,
                        val_pixels,
                        IMG_SIZE,
                        val_labels,
                        mode='val')    # 实例化数据生成器
    train_loader = paddle.io.DataLoader(train_dataset, batch_size=10, shuffle=True, num_workers=2)
    val_loader = paddle.io.DataLoader(val_dataset, batch_size=15, shuffle=False, num_workers=2)    # 实例化模型
    model = YOLOv3(num_classes = NUM_CLASSES)    
    # 学习率设置
    learning_rate = get_lr()    # 优化策略
    opt = paddle.optimizer.Momentum(
                 learning_rate=learning_rate,
                 momentum=0.9,
                 weight_decay=paddle.regularizer.L2Decay(0.0005),
                 parameters=model.parameters())  #创建优化器
    # opt = paddle.optimizer.Adam(learning_rate=learning_rate, weight_decay=paddle.regularizer.L2Decay(0.0005), parameters=model.parameters())
    
    # 开启训练
    for epoch in range(MAX_EPOCH):        for i, data in enumerate(train_loader()):
            img, gt_boxes, gt_labels, img_scale = data
            gt_scores = np.ones(gt_labels.shape).astype('float32')
            gt_scores = paddle.to_tensor(gt_scores)
            img = paddle.to_tensor(img)
            gt_boxes = paddle.to_tensor(gt_boxes)
            gt_labels = paddle.to_tensor(gt_labels)
            outputs = model(img)  # 前向传播，输出[P0, P1, P2]
            loss = model.get_loss(outputs, gt_boxes, gt_labels, gtscore=gt_scores,
                                  anchors = ANCHORS,
                                  anchor_masks = ANCHOR_MASKS,
                                  ignore_thresh=IGNORE_THRESH,
                                  use_label_smooth=False)  # 计算损失函数
            loss.backward()    # 反向传播计算梯度
            opt.step()  # 更新参数
            opt.clear_grad()            if i % 100 == 0:
                timestring = time.strftime("%Y-%m-%d %H:%M:%S",time.localtime(time.time()))                print('{}[TRAIN]epoch {}, iter {}, output loss: {}'.format(timestring, epoch, i, loss.numpy()))        # 保存参数
        if (epoch % 5 == 0) or (epoch == MAX_EPOCH -1) or (epoch == 6):  # 这里是为了保存第6轮
            paddle.save(model.state_dict(), 'yolo_epoch{}'.format(epoch))        # 每个epoch结束之后在验证集上进行测试
        model.eval()        for i, data in enumerate(val_loader()):
            img, gt_boxes, gt_labels, img_scale = data
            gt_scores = np.ones(gt_labels.shape).astype('float32')
            gt_scores = paddle.to_tensor(gt_scores)
            img = paddle.to_tensor(img)
            gt_boxes = paddle.to_tensor(gt_boxes)
            gt_labels = paddle.to_tensor(gt_labels)
            outputs = model(img)
            loss = model.get_loss(outputs, gt_boxes, gt_labels, gtscore=gt_scores,
                                  anchors = ANCHORS,
                                  anchor_masks = ANCHOR_MASKS,
                                  ignore_thresh=IGNORE_THRESH,
                                  use_label_smooth=False)            if i % 1 == 0:
                timestring = time.strftime("%Y-%m-%d %H:%M:%S",time.localtime(time.time()))                print('{}[VALID]epoch {}, iter {}, output loss: {}'.format(timestring, epoch, i, loss.numpy()))
        model.train()

from work.predict import predict

datadir = 'data/test'predict(datadir)

[[img_name, [[label, score, x1, y1, x2, y2], ..., [label, score, x1, y1, x2, y2]]], 
 [img_name, [[label, score, x1, y1, x2, y2], ..., [label, score, x1, y1, x2, y2]]],
  ...
 [img_name, [[label, score, x1, y1, x2, y2],..., [label, score, x1, y1, x2, y2]]]]

 [img_name, [[label, score, x1, y1, x2, y2],..., [label, score, x1, y1, x2, y2]]]

 [[label, score, x1, x2, y1, y2],..., [label, score, x1, y1, x2, y2]]

from work.fnc import fnfrom PIL import Image  
import matplotlib.pyplot as plt 


image_name = '2fd875eaa.jpg'plt.figure()
plt.imshow(fn(image_name, 0.15, 0.15))
plt.show()