說明:
本例程使用YOLOv3進行昆蟲檢測。例程分為數據處理、模型設計、損失函數、訓練模型、模型預測和測試模型六個部分。本篇為第一部分,實現了昆蟲檢測訓練的數據預處理功能和預測和測試時讀取和顯示數據功能。
數據集下載地址:https://aistudio.baidu.com/aistudio/datasetdetail/19748
實驗代碼:
數據增強:
import matplotlib.pyplot as plt import matplotlib.patches as patches from PIL import Image from source.data import get_data_list, get_data_item, augment_image # 讀取數據 train_set = './dataset/train/' data_list = get_data_list(train_set) # 讀取數據列表 data = data_list[0] # 讀取第一條數據 image, gtbox, gtcls, image_size = get_data_item(data) # 讀取數據項目 # 增強圖像 scale_size = (608, 608) # 縮放寬高 image, gtbox, gtcls = augment_image(image, gtbox, gtcls, scale_size) # 顯示圖像 object_names = ['Boerner', 'Leconte', 'Linnaeus', 'acuminatus', 'armandi', 'coleoptera', 'linnaeus'] # 物體名稱 color = ['r', 'g', 'b', 'c','m', 'y', 'k'] # 邊框顏色 image = Image.fromarray(image) # 轉換為Image格式 plt.figure(figsize=(10, 10)) # 設置顯示圖像大小 currentAxis = plt.gca() # 獲取圖像當前坐標 gtbox[:, 0] = gtbox[:, 0]*float(scale_size[1]) # 計算邊框x真實值 gtbox[:, 1] = gtbox[:, 1]*float(scale_size[0]) # 計算邊框y真實值 gtbox[:, 2] = gtbox[:, 2]*float(scale_size[1]) # 計算邊框w真實值 gtbox[:, 3] = gtbox[:, 3]*float(scale_size[0]) # 計算邊框h真實值 for i in range(len(gtbox)): # 繪制邊框 if gtbox[i, 2] > 1e-3 and gtbox[i, 3] > 1e-3: # 獲取數據 x = gtbox[i, 0] - gtbox[i, 2]/2 y = gtbox[i, 1] - gtbox[i, 3]/2 w = gtbox[i, 2] h = gtbox[i, 3] index = int(gtcls[i]) # 類別索引 names = object_names[index] # 類別名稱 # 繪制邊框 rectangle = patches.Rectangle((x, y), w, h, linewidth=1, edgecolor=color[index], facecolor=color[index], fill=False, linestyle='-') currentAxis.add_patch(rectangle) # 繪制類別 plt.text(x, y, names, fontsize=12, color=color[index]) plt.imshow(image) plt.show()
結果:
單線程讀取批次數據:
from source.data import single_thread_reader # 批次讀取訓練數據 train_set = './dataset/train/' train_reader = single_thread_reader(train_set, 2, 'train') train_data = next(train_reader()) print('train_data - image:{}, gtbox:{}, gtcls:{}, image_size:{}'.format( train_data[0].shape, train_data[1].shape, train_data[2].shape, train_data[3].shape))
結果:train_data - image:(2, 3, 448, 448), gtbox:(2, 50, 4), gtcls:(2, 50), image_size:(2, 2)
多線程讀取批次數據:
from source.data import multip_thread_reader # 多線程讀取訓練數據 valid_set = './dataset/val/' valid_reader = multip_thread_reader(valid_set, 2, 'valid') valid_data = next(valid_reader()) print('valid_data - image:{}, gtbox:{}, gtcls:{}, image_size:{}'.format( valid_data[0].shape, valid_data[1].shape, valid_data[2].shape, valid_data[3].shape))
結果:valid_data - image:(2, 3, 512, 512), gtbox:(2, 50, 4), gtcls:(2, 50), image_size:(2, 2)
data.py文件
import os import random import numpy as np import xml.etree.ElementTree as ET from PIL import Image, ImageEnhance import matplotlib.pyplot as plt import matplotlib.patches as patches import paddle object_names = ['Boerner', 'Leconte', 'Linnaeus', 'acuminatus', 'armandi', 'coleoptera', 'linnaeus'] # 物體名稱 def get_object_gtcls(): """ 功能: 將物體名稱映射成物體類別 輸入: 輸出: object_gtcls - 物體類別 """ object_gtcls = {} # 物體類別字典 for key, value in enumerate(object_names): object_gtcls[value] = key # 將物體名稱映射成物體類別 return object_gtcls def get_data_list(data_path): """ 功能: 讀取數據列表 輸入: data_path - 數據路徑 輸出: data_list - 數據列表 """ file_list = os.listdir(os.path.join(data_path, 'annotations', 'xmls')) # 文件列表 data_list = [] # 數據列表 for file_id, file_name in enumerate(file_list): # 讀取數據文件 file_path = os.path.join(data_path, 'annotations', 'xmls', file_name) # 文件路徑 xml_tree = ET.parse(file_path) # 解析文件 # 讀取圖像數據 image_path = os.path.join(data_path, 'images', file_name.split('.')[0] + '.jpeg') # 圖像路徑 image_w = float(xml_tree.find('size').find('width').text) # 圖像寬度 image_h = float(xml_tree.find('size').find('height').text) # 圖像高度 # 讀取物體數據 object_list = xml_tree.findall('object') # 物體列表 object_gtbox = np.zeros((len(object_list), 4), dtype=np.float32) # 物體邊框 object_gtcls = np.zeros((len(object_list), ), dtype=np.int32) # 物體類別 for object_id, object_item in enumerate(object_list): # 讀取物體類別 object_name = object_item.find('name').text # 讀取物體名稱 object_gtcls[object_id] = get_object_gtcls()[object_name] # 將物體名稱映射成物體類別 # 讀取物體邊框 x_min = float(object_item.find('bndbox').find('xmin').text) y_min = float(object_item.find('bndbox').find('ymin').text) x_max = float(object_item.find('bndbox').find('xmax').text) y_max = float(object_item.find('bndbox').find('ymax').text) x_min = max(0.0, x_min) y_min = max(0.0, y_min) x_max = min(x_max, image_w - 1.0) y_max = min(y_max, image_h - 1.0) object_gtbox[object_id] = [ (x_min + x_max) / 2.0, (y_min + y_max) / 2.0, (x_max - x_min + 1.0), (y_max - y_min + 1.0)] # 計算物體邊框: xywh格式 # 保存數據列表 data = {'image_path': image_path, 'image_width': image_w, 'image_height': image_h, 'object_gtbox': object_gtbox, 'object_gtcls': object_gtcls} if len(data) != 0: data_list.append(data) return data_list ############################################################################################################## def get_object_item(object_gtbox, object_gtcls, image_size): """ 功能: 讀取物體項目 輸入: object_gtbox - 物體邊框 object_gtcls - 物體類別 image_size - 圖像高寬 輸出: gtbox - 邊框列表 gtcls - 類別列表 """ # 添加物體數據列表 max_num = 50 # 最大物體數量 gtbox = np.zeros((max_num, 4)) # 邊框列表 gtcls = np.zeros((max_num, )) # 類別列表 for i in range(len(object_gtbox)): gtbox[i, :] = object_gtbox[i, :] gtcls[i] = object_gtcls[i] if i > max_num: # 是否超過最大物體數量 break # 轉換為相對真實圖像比例的邊框 gtbox[:, 0] = gtbox[:, 0]/float(image_size[1]) # 計算邊框x相對值 gtbox[:, 1] = gtbox[:, 1]/float(image_size[0]) # 計算邊框y相對值 gtbox[:, 2] = gtbox[:, 2]/float(image_size[1]) # 計算邊框w相對值 gtbox[:, 3] = gtbox[:, 3]/float(image_size[0]) # 計算邊框h相對值 return gtbox, gtcls def get_data_item(data): """ 功能: 讀取數據項目 輸入: data - 數據 輸出: image - 圖像數據 gtbox - 邊框列表 gtcls - 類別列表 image_size - 圖像高寬 """ # 讀取數據項目 image_path = data['image_path'] # 圖像路徑 image_w = data['image_width'] # 圖像寬度 image_h = data['image_height'] # 圖像高度 object_gtbox = data['object_gtbox'] # 物體邊框 object_gtcls = data['object_gtcls'] # 物體類別 # 打開圖像文件 image = Image.open(image_path) # 打開圖像 if image.mode != 'RGB': image = image.convert('RGB') image = np.array(image) # 轉換為ndarray格式, 讀取數據類型為HWC,uint8類型 # 檢查圖像高寬 assert image.shape[0] == int(image_h), \ 'image path: {}, image.shape[0]:{} != image_h: {}'.format( image_path, image.shape[0], image_h) assert image.shape[1] == int(image_w), \ 'image path: {}, image.shape[1]:{} != image_w: {}'.format( image_path, image.shape[1], image_w) # 讀取物體項目 image_size = (image_h, image_w) gtbox, gtcls = get_object_item(object_gtbox, object_gtcls, image_size) return image, gtbox, gtcls, image_size ############################################################################################################## def random_distort_image(image): """ 功能: 隨機變換圖像 輸入: image - 圖像數據 輸出: image - 圖像數據 """ # 隨機變換飽和度 def random_distort_saturation(image): if random.random() > 0.5: delta = random.uniform(-0.5, 0.5) + 1 image = ImageEnhance.Color(image).enhance(delta) return image # 隨機變換明亮度 def random_distort_brightness(image): if random.random() > 0.5: delta = random.uniform(-0.5, 0.5) + 1 image = ImageEnhance.Brightness(image).enhance(delta) return image # 隨機變換對比度 def random_distort_constract(image): if random.random() > 0.5: delta = random.uniform(-0.5, 0.5) + 1 image = ImageEnhance.Contrast(image).enhance(delta) return image # 隨機變換色調 def random_distort_hue(image): if random.random() > 0.5: delta = random.uniform(-18, 18) image = np.array(image.convert('HSV')) image[:, :, 0] = image[:, :, 0] + delta # 變換色調 image = Image.fromarray(image, mode='HSV').convert('RGB') return image # 隨機變換圖像 distort = [random_distort_saturation, random_distort_brightness, random_distort_constract, random_distort_hue] # 變換方法列表 np.random.shuffle(distort) # 打亂變換順序 image = Image.fromarray(image) # 轉換為Image格式 image = distort[0](image) image = distort[1](image) image = distort[2](image) image = distort[3](image) image = np.asarray(image) # 轉換為ndarray格式 return image def random_expand_image(image, gtbox, keep_ratio=True): """ 功能: 隨機填充圖像 輸入: image - 圖像數據 gtbox - 邊框列表 keep_ratio - 保持寬高比例 輸出: image - 圖像數據 gtbox - 邊框列表 """ # 是否填充圖像 if random.random() > 0.5: return image, gtbox # 生成填充比例 max_ratio = 4 x_ratio = random.uniform(1, max_ratio) # 隨機產生x填充比例 if keep_ratio: # 是否保持寬高比例 y_ratio = x_ratio # y填充比例等於x填充比例 else: y_ratio = random.uniform(1, max_ratio) # 隨機產生y填充比例 # 計算填充寬高 image_h, image_w, image_channel = image.shape # 獲取圖像高寬和通道數 expand_w = int(image_w * x_ratio) # 計算填充寬度 expand_h = int(image_h * y_ratio) # 計算填充高度 x_offset = random.randint(0, expand_w - image_w) # 隨機生成原圖在填充圖中x坐標位置 y_offset = random.randint(0, expand_h - image_h) # 隨機生成原圖在填充圖中y坐標位置 # 生成填充圖像 expand_image = np.zeros((expand_h, expand_w, image_channel)) # float32類型 mean_value = [0.485, 0.456, 0.406] # COCO數據集通道平均值 for i in range(image_channel): expand_image[:, :, i] = mean_value[i] * 255.0 # 使用均值填充每個通道 # 填充原始圖像 expand_image[y_offset : y_offset + image_h, x_offset : x_offset + image_w, :] = image # 填充原始圖像 image = expand_image.astype('uint8') # 轉換為uint8類型 # 計算相對邊框 gtbox[:, 0] = ((gtbox[:, 0] * image_w) + x_offset) / float(expand_w) # 計算邊框x相對值 gtbox[:, 1] = ((gtbox[:, 1] * image_h) + y_offset) / float(expand_h) # 計算邊框y相對值 gtbox[:, 2] = gtbox[:, 2] / x_ratio # 計算邊框w相對值 gtbox[:, 3] = gtbox[:, 3] / y_ratio # 計算邊框h相對值 return image, gtbox def get_boxes_ious_xywh(box1, box2): """ 功能: 計算邊框列表的交並比 輸入: box1 - 邊框列表1 box2 - 邊框列表2 輸出: ious - 交並比值列表 """ # 判斷邊框維度 assert box1.shape[-1] == 4, "Box1 shape[-1] should be 4." assert box2.shape[-1] == 4, "Box2 shape[-1] should be 4." # 計算交集面積 x1_min = box1[:, 0] - box1[:, 2]/2.0 y1_min = box1[:, 1] - box1[:, 3]/2.0 x1_max = box1[:, 0] + box1[:, 2]/2.0 y1_max = box1[:, 1] + box1[:, 3]/2.0 x2_min = box2[:, 0] - box2[:, 2]/2.0 y2_min = box2[:, 1] - box2[:, 3]/2.0 x2_max = box2[:, 0] + box2[:, 2]/2.0 y2_max = box2[:, 1] + box2[:, 3]/2.0 x_min = np.maximum(x1_min, x2_min) y_min = np.maximum(y1_min, y2_min) x_max = np.minimum(x1_max, x2_max) y_max = np.minimum(y1_max, y2_max) w = np.maximum(x_max - x_min, 0.0) h = np.maximum(y_max - y_min, 0.0) intersection = w * h # 交集面積 # 計算並集面積 s1 = box1[:, 2] * box1[:, 3] s2 = box2[:, 2] * box2[:, 3] union = s1 + s2 - intersection # 並集面積 # 計算交並比值 ious = intersection / union return ious def get_cpbox_item(cpbox, gtbox, gtcls, image_size): """ 功能: 計算裁剪邊框 輸入: cpbox - 裁剪邊框: 真實位置 gtbox - 真實邊框: 相對位置 gtcls - 真實類別 image_size - 圖像尺寸 輸出: gtbox - 裁剪邊框 gtcls - 裁剪類別 box_number - 邊框數量 """ # 拷貝真實邊框列表 gtbox = gtbox.copy() # 防止本次返回失敗, 影響下次真實邊框值 gtcls = gtcls.copy() # 防止本次返回失敗, 影響下次真實類別值 # 轉換真實邊框位置:x1, y1, x2, y2 image_w, image_h = map(float, image_size) # 讀取圖像寬高 gtbox[:, 0], gtbox[:, 2] = \ (gtbox[:, 0] - gtbox[:, 2]/2) * image_w, (gtbox[:, 0] + gtbox[:, 2]/2) * image_w # 先計算右邊,再賦值左邊,避免影響左邊原始值 gtbox[:, 1], gtbox[:, 3] = \ (gtbox[:, 1] - gtbox[:, 3]/2) * image_h, (gtbox[:, 1] + gtbox[:, 3]/2) * image_h # 原址計算避免內存拷貝,分開計算,需要先拷貝gtbox # 轉換裁剪邊框位置 cpbox_x, cpbox_y, cpbox_w, cpbox_h = map(float, cpbox) # 讀取邊框位置 cpbox = np.array([cpbox_x, cpbox_y, cpbox_x + cpbox_w, cpbox_y + cpbox_h]) # 轉換為ndarray格式:xyxy格式 # 計算真實邊框中心在裁剪邊框內的個數的掩碼 gtbox_centers = (gtbox[:, :2] + gtbox[:, 2:]) / 2.0 # 計算所有真實邊框的中心位置: x=(x1+x2)/2, y=(y1+y2)/2 mask = np.logical_and(cpbox[:2] <= gtbox_centers, gtbox_centers <= cpbox[2:]).all(axis=1) # 中心位置掩碼:真實邊框中心是否在裁剪邊框內 # 計算裁剪邊框位置 gtbox[:, :2] = np.maximum(gtbox[:, :2], cpbox[:2]) # 計算真實邊框與裁剪邊框最大x1y1值 gtbox[:, 2:] = np.minimum(gtbox[:, 2:], cpbox[2:]) # 計算真實邊框與裁剪邊框最小x2y2值 gtbox[:, :2] -= cpbox[:2] # 計算真實邊框x1y1值, 相對於裁剪邊框的左上角位置 gtbox[:, 2:] -= cpbox[:2] # 計算真實邊框x2y2值, 相對於裁剪邊框的左上角位置 # 計算裁剪邊框中左上角位置小於右下角的掩碼 mask = np.logical_and(mask, (gtbox[:, :2] < gtbox[:, 2:]).all(axis=1) ) # 裁剪邊框的右下坐標是否在左上角左邊之下 # 計算裁剪邊框位置 gtbox = gtbox * np.expand_dims(mask.astype('float32'), axis=1) # 使用掩碼計算合格的裁剪邊框 # 轉換相對邊框位置: x, y, w, h gtbox[:, 0], gtbox[:, 2] = \ (gtbox[:, 0] + gtbox[:, 2])/2 / cpbox_w, (gtbox[:, 2] - gtbox[:, 0]) / cpbox_w # 先計算右邊,再賦值左邊,避免影響左邊原始值 gtbox[:, 1], gtbox[:, 3] = \ (gtbox[:, 1] + gtbox[:, 3])/2 / cpbox_h, (gtbox[:, 3] - gtbox[:, 1]) / cpbox_h # 原址計算避免內存拷貝,分開計算,需要先拷貝gtbox # 計算裁剪邊框類別 gtcls = gtcls * mask.astype('float32') # 使用掩碼計算合格的裁剪邊框 # 計算裁剪邊框數量 box_number = mask.sum() # 統計掩碼為1的數量 return gtbox, gtcls, box_number def random_crop_image(image, gtbox, gtcls, crop_scope=[0.3, 1.0], max_ratio=2.0, iou_scopes=None, max_trial=50): """ 功能: 隨機裁剪圖像 輸入: image - 圖像數據 gtbox - 邊框列表 gtcls - 類別列表 crop_scope - 裁剪比例范圍 max_ratio - 最大裁剪比例 iou_scopes - 交並比值范圍 max_trial - 最大試驗次數 輸出: image - 圖像數據 gtbox - 邊框列表 gtcls - 類別列表 """ # 是否存在物體 if random.random() > 0.5: return image, gtbox, gtcls if len(gtbox) == 0: # 如果物體邊框數為零, 則返回原始圖像數據項目 return image, gtbox, gtcls # 交並比值范圍 if not iou_scopes: # 如果不存在交並比值范圍列表, 則使用交默認並比值范圍列表 iou_scopes = [(0.1, 1.0), (0.3, 1.0), (0.5, 1.0), (0.7, 1.0), (0.9, 1.0), (0.0, 1.0)] # 轉換圖像格式 image = Image.fromarray(image) # 轉換為Image格式 image_w, image_h = image.size # 獲取圖像寬高 # 計算裁剪邊框 cpbox_list = [(0, 0, image_w, image_h)] # 裁剪邊框列表 for min_iou, max_iou in iou_scopes: for i in range(max_trial): # 隨機生成裁剪比例 crop_ratio = random.uniform(crop_scope[0], crop_scope[1]) # 隨機生成裁剪寬高比例 aspect_ratio = random.uniform( max(1 / max_ratio, (crop_ratio * crop_ratio)), min(max_ratio, 1 / (crop_ratio * crop_ratio))) # 隨機生成裁剪縱橫比例 # 計算裁剪邊框位置 crop_h = int(image_h * crop_ratio / np.sqrt(aspect_ratio)) # 計算隨機生成的裁剪高度 crop_w = int(image_w * crop_ratio * np.sqrt(aspect_ratio)) # 計算隨機生成的裁剪寬度 crop_x = random.randint(0, image_w - crop_w) # 隨機生成裁剪x坐標 crop_y = random.randint(0, image_h - crop_h) # 隨機生成裁剪y坐標 # 計算邊框交並集值 cpbox = np.array( [[(crop_x + crop_w / 2.0) / float(image_w), (crop_y + crop_h / 2.0) / float(image_h), crop_w / float(image_w), crop_h / float(image_h)]]) # 計算裁剪邊框相對位置: xywh格式 ious = get_boxes_ious_xywh(cpbox, gtbox) # 裁剪邊框形狀為(1, 4), 真實邊框形狀為(50, 4) # 添加裁剪邊框列表 if min_iou <= ious.min() and ious.max() <= max_iou: # 如果符合交並比值范圍, 則添加一個裁剪邊框, 並結束循環 cpbox_list.append((crop_x, crop_y, crop_w, crop_h)) # 裁剪邊框為真實位置 break # 隨機裁剪圖像 for i in range(len(cpbox_list)): # 彈出裁剪邊框 cpbox = cpbox_list.pop(random.randint(0, len(cpbox_list) - 1)) # 隨機彈出裁剪邊框 # 計算裁剪邊框 crop_boxes, crop_gtcls, box_number = get_cpbox_item(cpbox, gtbox, gtcls, image.size) # 開始裁剪圖像 if box_number > 0: # 如果裁剪邊框數量大於0, 則裁剪圖像,並結束循環 image = image.crop( (cpbox[0], cpbox[1], cpbox[0] + cpbox[2], cpbox[1] + cpbox[3]) ) # 用真實位置裁剪圖像 image = image.resize(image.size, Image.LANCZOS) # 高質量縮放到原來大小 gtbox = crop_boxes # 返回裁剪邊框 gtcls = crop_gtcls # 返回裁剪類別 break # 轉換圖像格式 image = np.asarray(image) # 轉換為ndarray格式 return image, gtbox, gtcls def random_interpolate_image(image, scale_size, interpolation=None): """ 功能: 隨機插值圖像 輸入: image - 圖像數據 scale_size - 縮放寬高 interpolation - 插值方法 輸出: image - 圖像數據 """ # 轉換圖像格式 image = Image.fromarray(image) # 轉換為Image格式 # 隨機縮放圖像 interpolation_method = [Image.NEAREST,Image.BILINEAR ,Image.BICUBIC, Image.LANCZOS] # 插值方法列表 if not interpolation or interpolation not in interpolation_method: interpolation = interpolation_method[random.randint(0, len(interpolation_method) - 1)] # 隨機選取插值方法 image = image.resize(scale_size, interpolation) # 轉換圖像格式 image = np.asarray(image) # 轉換為ndarray格式,數據類型為HWC,uint8類型 return image def random_flip_image(image, gtbox): """ 功能: 隨機翻轉圖像 輸入: image - 圖像數據 gtbox - 邊框列表 輸出: image - 圖像數據 gtbox - 邊框列表 """ if random.random() > 0.5: image = image[:, ::-1, :] # 水平翻轉圖像, 列全部倒序排列 gtbox[:, 0] = 1.0 - gtbox[:, 0] # 水平翻轉邊框, x坐標變為相反數 return image, gtbox def random_shuffle_gtbox(gtbox, gtcls): """ 功能: 隨機打亂邊框 輸入: gtbox - 邊框列表 gtcls - 類別列表 輸出: gtbox - 邊框列表 gtcls - 類別列表 """ # 連接邊框和類別 data_list = np.concatenate([gtbox, gtcls[:, np.newaxis]], axis=1) # 打亂列表的順序 index = np.arange(data_list.shape[0]) np.random.shuffle(index) data_list = data_list[index, :] # 保存邊框和類別 gtbox = data_list[:, :4] gtcls = data_list[:, -1] return gtbox, gtcls def augment_image(image, gtbox, gtcls, scale_size): """ 功能: 增強圖像 輸入: image - 圖像數據 gtbox - 邊框列表 gtcls - 類別列表 scale_size - 縮放尺寸 輸出: image - 圖像數據 gtbox - 邊框列表 gtcls - 類別列表 """ # 隨機變換圖像 image = random_distort_image(image) # 隨機填充圖像 image, gtbox = random_expand_image(image, gtbox) # 隨機裁剪圖像 image, gtbox, gtcls = random_crop_image(image, gtbox, gtcls) # 隨機插值圖像 image = random_interpolate_image(image, scale_size) # 隨機翻轉圖像 image, gtbox = random_flip_image(image, gtbox) # 隨機打亂邊框 gtbox, gtcls = random_shuffle_gtbox(gtbox, gtcls) return image, gtbox, gtcls ############################################################################################################## def get_scale_size(mode): """ 功能: 獲取縮放尺寸 輸入: mode - 獲取模式 輸出: scale_size - 縮放尺寸 """ if (mode == 'train') or (mode == 'valid'): # 如果是訓練或驗證模式, 隨機生成縮放寬高 scale_size = 320 + 32 * random.randint(0, 9) # 隨機生成寬高,范圍為[320, 608],步進為32 else: scale_size = 608 scale_size = (scale_size, scale_size) # 組合縮放寬高[w,h] return scale_size def get_data_array(batch_data): """ 功能: 將數據列表轉換為數組構成的元組 輸入: batch_data - 批次數據列表 輸出: image_array - 圖像數據數組 gtbox_array - 物體邊框數組 gtcls_array - 圖像類別數組 image_size_array - 圖像寬高數組 """ image_array = np.array([item[0] for item in batch_data], dtype='float32') gtbox_array = np.array([item[1] for item in batch_data], dtype='float32') gtcls_array = np.array([item[2] for item in batch_data], dtype='int32') image_size_array = np.array([item[3] for item in batch_data], dtype='int32') return image_array, gtbox_array, gtcls_array, image_size_array def data_reader(data, scale_size): """ 功能: 讀取一條數據 輸入: data - 數據項目 scale_size - 縮放圖像寬高 輸出: image - 圖像數據 gtbox - 邊框列表 gtcls - 類別列表 image_size - 原始圖像高寬 """ # 讀取數據 image, gtbox, gtcls, image_size = get_data_item(data) # 增強圖像 image, gtbox, gtcls = augment_image(image, gtbox, gtcls, scale_size) # 減去均值 mean = np.array([0.485, 0.456, 0.406]).reshape((1, 1, -1)) # COCO數據集通道平均值 stdv = np.array([0.229, 0.224, 0.225]).reshape((1, 1, -1)) # COCO數據集通道標准差 image = (image/255.0 - mean) / stdv # 對圖像進行歸一化 image = image.astype('float32').transpose((2, 0, 1)) # 轉換圖片格式:[H,W,C]到[C,H,W] return image, gtbox, gtcls, image_size def single_thread_reader(data_path, batch_size=8, mode='train'): """ 功能: 單線程讀取批次數據 輸入: data_path - 數據集路徑 batch_size - 每批數據大小 mode - 讀取數據模式: train或valid 輸出: reader - 數據讀取器 """ # 讀取數據列表 data_list = get_data_list(data_path) # 讀取數據項目 def reader(): # 設置讀取模式 if mode == 'train': # 如果是訓練模式, 則打亂數據列表 np.random.shuffle(data_list) scale_size = get_scale_size(mode) # 設置縮放寬高 # 輸出批次數據 batch_data = [] # 批次數據 for item in data_list: # 讀取數據 image, gtbox, gtcls, image_size = data_reader(item, scale_size) # 輸出數據 batch_data.append((image, gtbox, gtcls, image_size)) if len(batch_data) == batch_size: # 如果壓入一批數據, 則彈出數據 # 彈出數據 yield get_data_array(batch_data) # 重置數據 batch_data = [] image_size = get_scale_size(mode) # 輸出剩余數據 if len(batch_data) > 0: yield get_data_array(batch_data) return reader def multip_thread_reader(data_path, batch_size=8, mode='train'): """ 功能: 多線程讀取批次數據 輸入: data_path - 數據集路徑 batch_size - 每批數據大小 mode - 讀取數據模式: train 或 valid 輸出: reader - 數據讀取器 """ # 讀取數據列表 data_list = get_data_list(data_path) # 讀取數據項目 def item_loader(): # 設置讀取模式 if mode == 'train': # 如果是訓練模式, 則打亂數據列表 np.random.shuffle(data_list) scale_size = get_scale_size(mode) # 設置縮放寬高 # 輸出批次項目 batch_item = [] # 批次項目 for item in data_list: # 輸出項目 batch_item.append((item, scale_size)) if len(batch_item) == batch_size: # 如果壓入一批項目, 則彈出項目 # 彈出數據 yield batch_item # 重置項目 batch_item = [] image_size = get_scale_size(mode) # 輸出剩余數據 if len(batch_item) > 0: yield batch_item # 讀取數據內容 def data_loader(batch_item): batch_data = [] # 批次數據 for item, scale_size in batch_item: image, gtbox, label, image_size = data_reader(item, scale_size) batch_data.append((image, gtbox, label, image_size)) return get_data_array(batch_data) # 多線程讀取器 reader = paddle.reader.xmap_readers(data_loader, item_loader, process_num=4, buffer_size=16) return reader ############################################################################################################## def single_test_reader(image_path, scale_size=(608, 608)): """ 功能: 讀取一張預測圖像 輸入: image_path - 圖像路徑 scale_size - 縮放高寬 輸出: image - 圖像數據 image_size - 原始高寬 """ # 讀取圖像 image = Image.open(image_path) # 讀取圖像 if image.mode != 'RGB': image = image.convert('RGB') image_size = (image.size[0], image.size[1]) # 讀取尺寸 image = image.resize(scale_size, Image.BILINEAR) # 縮放圖像 # 轉換格式 image = np.array(image, dtype='float32') image_size = np.array(image_size, dtype='int32') # 減去均值 mean = np.array([0.485, 0.456, 0.406]).reshape((1, 1, -1)) # COCO數據集通道平均值 stdv = np.array([0.229, 0.224, 0.225]).reshape((1, 1, -1)) # COCO數據集通道標准差 image = (image/255.0 - mean) / stdv # 對圖像進行歸一化 image = image.astype('float32').transpose((2, 0, 1)) # 轉換圖片格式:[H,W,C]到[C,H,W] # 增加維度 image = np.expand_dims(image, axis=0) # 增加維度: (1,3,608,608) image_size = np.expand_dims(image_size, axis=0) # 增加維度: (1,2) return image, image_size def display_infer(infer, image_path): """ 功能: 顯示預測結果 輸入: infer - 預測結果 image_path - 圖像路徑 輸出: """ # 讀取圖像 image = Image.open(image_path) # 讀取圖像 if image.mode != 'RGB': image = image.convert('RGB') # 轉換格式 # 繪制結果 object_names = ['Boerner','Leconte','Linnaeus','acuminatus','armandi','coleoptera','linnaeus'] # 物體名稱 color = ['r', 'g', 'b', 'c','m', 'y', 'k'] # 邊框顏色 plt.figure(figsize=(10, 10)) # 設置顯示圖像大小 currentAxis = plt.gca() # 獲取圖像當前坐標 for item in infer: # 遍歷預測結果 # 獲取結果 index = int(item[0]) # 類別索引 names = object_names[index] # 類別名稱 pdbox = item[2:6] # 邊框位置 # 繪制邊框 rectangle = patches.Rectangle( # 設置邊框 (pdbox[0], pdbox[1]), pdbox[2]-pdbox[0]+1, pdbox[3]-pdbox[1]+1, linewidth=1, edgecolor=color[index], facecolor=color[index], fill=False, linestyle='-') currentAxis.add_patch(rectangle) # 繪制邊框 plt.text(pdbox[0], pdbox[1], names, fontsize=12, color=color[index]) # 繪制類別 # 顯示結果 plt.imshow(image) plt.show() def get_test_array(batch_data): """ 功能: 將數據列表轉換為數組構成的元組 輸入: batch_data - 數據列表 輸出: image_name_array - 圖像名字數組 image_array - 圖像數據數組 image_size_array - 圖像寬高數組 """ image_name_array = np.array([item[0] for item in batch_data]) image_array = np.array([item[1] for item in batch_data], dtype='float32') image_size_array = np.array([item[2] for item in batch_data], dtype='int32') return image_name_array, image_array, image_size_array def multip_test_reader(data_path, batch_size=1, scale_size=(608, 608)): """ 功能: 讀取一批測試數據 輸入: data_path - 數據目錄 batch_size - 每批大小 scale_size - 縮放高寬 輸出: image_name - 圖像名稱 image - 圖像數據 image_size - 原始高寬 """ # 讀取數據列表 data_list = os.listdir(data_path) # 讀取一批數據 def reader(): # 輸出批次數據 batch_data = [] # 批次數據 for image_name in data_list: # 讀取圖像路徑 image_path = os.path.join(data_path, image_name) # 讀取一張圖像 image = Image.open(image_path) # 讀取圖像 if image.mode != 'RGB': image = image.convert('RGB') image_size = (image.size[0], image.size[1]) # 圖像大小 image = image.resize(scale_size, Image.BILINEAR) # 縮放圖像 # 轉換格式 image = np.array(image, dtype='float32') image_size = np.array(image_size, dtype='int32') # 減去均值 mean = np.array([0.485, 0.456, 0.406]).reshape((1, 1, -1)) # COCO數據集通道平均值 stdv = np.array([0.229, 0.224, 0.225]).reshape((1, 1, -1)) # COCO數據集通道標准差 image = (image/255.0 - mean) / stdv # 對圖像進行歸一化 image = image.astype('float32').transpose((2, 0, 1)) # 轉換圖片格式:[H,W,C]到[C,H,W] # 輸出數據 batch_data.append((image_name.split('.')[0], image, image_size)) if len(batch_data) == batch_size: # 如果壓入一批數據, 則彈出數據 # 彈出數據 yield get_test_array(batch_data) # 重置數據 batch_data = [] # 輸出剩余數據 if len(batch_data) > 0: yield get_test_array(batch_data)
參考資料:
https://aistudio.baidu.com/aistudio/projectdetail/742781
https://aistudio.baidu.com/aistudio/projectdetail/672017
https://aistudio.baidu.com/aistudio/projectdetail/868589
https://aistudio.baidu.com/aistudio/projectdetail/122277