1.背景
在目標檢測中,需要進行圖像增強。這里的代碼模擬了旋轉、扭曲圖像的功能,并且在扭曲的時候,能夠同時把標注的結果也進行扭曲。
這里忽略了讀取xml的過程,假設圖像IMG存在對應的標注框,且坐標為左上、右下兩個點,這立刻喲把IMG進行扭曲,并且盡可能保證標簽的坐標點也隨之扭曲。
2.效果
我有一張鐵塔的照片(原圖略,原圖是一張角度非常正確的圖片),處理后效果如下圖所示,可以看到圖像出現了扭曲(出現了黑邊),并且圖像右側中間部分的紅色絕緣子旁邊出現了兩個小白點(為了便于觀察而添加的),說明起到了扭曲的作用,并且大概率能保證坐標也被正確變化。
3.動機
設計這種數據增強策略,主要出于以下兩種考慮:
- 在實際使用中,我發現即便是yolov8這樣的模型,對于扭曲的圖像(例如扭曲的筆記本),識別效果會有非常嚴重的下降,而真實情況中非常有可能出現扭曲(例如有人把筆記本半合上),模型需要識別。在電力場景中,也可能出現拍攝角度異常等情況。
- 在訓練中,將不同角度的目標展現給模型,也可能提高模型的效果。例如上圖中的絕緣子,比原圖要“矮胖”一些,因此可以算作“新樣本”,為模型提供更豐富的數據來源。
4.具體實現
第一步是圖像扭曲。在這一步中,代碼里的trapezoid_vertices表示你希望把圖像的哪一部分進行扭曲,target_trapezoid_vertices表示你希望把trapezoid_vertices扭曲到什么位置。你可以把代碼中相關位置的generate_random_cut和generate_random_perspective都去掉,然后就能明白了。
第二步是轉化原始標注的框。在使用gpt生成代碼的時候,提示詞將任務劃分為了5個步驟,提示詞如下
假設原始圖像PIC寬度為W,高度為H。在圖像標注場景中,我會告訴你兩個點special_points,分別代表原始的標注框的左上角和右下角。你需要執行以下步驟:
【步驟1】根據special_points,恢復原始的標注框的四個坐標。注意保留這四個點的順序關系。
【步驟2】:利用自定義函數fun1計算出變化之后的四個點的新坐標。
【步驟3】依次判斷四個點所連成的線段是否與原始圖像PIC的邊緣有交點,如果有,就將這些交點的坐標存到point_list中。
【步驟4】判斷變換后的四個點有哪些位于原始圖像PIC的范圍內,如果有,則將這些點的坐標存到point_list中。
【步驟5】分析point_list中的所有點,找出一個能完整包含這些點的矩形R,返回這個矩形R的左上角和右下角。
輔助代碼包含了前面提到的5個步驟,具體如下:
import random
import cv2
import numpy as np# 原始圖像剪裁的范圍
max_cut = 0.1
min_cut = 0.05# 扭曲的最大和最小程度
max_perspective = 0.1
min_perspective = 0.05# 制定初始剪裁范圍,你可以選擇從某些地方開始剪裁你的圖像
def generate_random_cut(x):# 計算 x 的 10% 和 20%min_value = min_perspective * xmax_value = max_perspective * x# 生成一個在 10% 到 20% 范圍內的隨機數random_value = random.uniform(min_value, max_value)# 隨機決定這個數是正數還是負數sign = random.choice([-1, 1])return sign * random_value# 制定目標梯形的范圍
def generate_random_perspective(x):# 計算 x 的 10% 和 20%min_value = min_perspective * xmax_value = max_perspective * xrandom_value = random.uniform(min_value, max_value)# 隨機決定這個數是正數還是負數sign = random.choice([-1, 1])return sign * random_value# 根據輸入的兩個點的坐標,復原出連續的矩形框的四個點坐標
def get_bbox_corners(special_points):top_left = special_points[0][0]bottom_right = special_points[0][1]top_right = [bottom_right[0], top_left[1]]bottom_left = [top_left[0], bottom_right[1]]return np.array([top_left, top_right, bottom_right, bottom_left])# line_intersection所需要的輔助函數
def on_segment(p, q, r):if (q[0] <= max(p[0], r[0]) and q[0] >= min(p[0], r[0]) andq[1] <= max(p[1], r[1]) and q[1] >= min(p[1], r[1])):return Truereturn False# line_intersection所需要的輔助函數
def orientation(p, q, r):val = (q[1] - p[1]) * (r[0] - q[0]) - (q[0] - p[0]) * (r[1] - q[1])if val == 0:return 0 # collinearelif val > 0:return 1 # clockwiseelse:return 2 # counterclockwise# line_intersection所需要的輔助函數
def segments_intersect(p1, q1, p2, q2):o1 = orientation(p1, q1, p2)o2 = orientation(p1, q1, q2)o3 = orientation(p2, q2, p1)o4 = orientation(p2, q2, q1)if o1 != o2 and o3 != o4:return Trueif o1 == 0 and on_segment(p1, p2, q1):return Trueif o2 == 0 and on_segment(p1, q2, q1):return Trueif o3 == 0 and on_segment(p2, p1, q2):return Trueif o4 == 0 and on_segment(p2, q1, q2):return Truereturn False# 判斷四個點所連成的線段是否與原始圖像的邊緣有交點
def line_intersection(p1, p2, edge_start, edge_end):if not segments_intersect(p1, p2, edge_start, edge_end):return Nonexdiff = (p1[0] - p2[0] + 0.01, edge_start[0] - edge_end[0] + 0.01)ydiff = (p1[1] - p2[1] + 0.01, edge_start[1] - edge_end[1] + 0.01)def det(a, b):return a[0] * b[1] - a[1] * b[0]div = det(xdiff, ydiff)if div == 0:return Noned = (det(p1, p2), det(edge_start, edge_end))x = det(d, xdiff) / divy = det(d, ydiff) / divreturn x, y# 判斷四個點所連成的線段是否與原始圖像的邊緣有交點
def check_intersections(bbox_corners, w, h):edges = [((0.01, 0), (w, 0.01)),((w, 0.01), (w, h)),((w, h), (0.01, h)),((0.01, h), (0.01, 0.01))]point_list = []for i in range(len(bbox_corners)):p1 = bbox_corners[i]p2 = bbox_corners[(i + 1) % len(bbox_corners)]for edge in edges:intersection = line_intersection(p1, p2, edge[0], edge[1])if intersection:point_list.append(intersection)return point_list# 判斷變換后的四個點是否位于原始圖像范圍內
def check_points_within_image(bbox_corners, w, h):point_list = []for point in bbox_corners:if 0 <= point[0] <= w and 0 <= point[1] <= h:point_list.append(point)return point_list# 分析所有點,找出一個能完整包含這些點的矩形
def get_bounding_box(point_list):min_x = min(point[0] for point in point_list)max_x = max(point[0] for point in point_list)min_y = min(point[1] for point in point_list)max_y = max(point[1] for point in point_list)return (min_x, min_y), (max_x, max_y)
使用方法如下:
# 數據增強 - 拉伸
if __name__ == "__main__":# 定義兩個特殊點,即原始標注框的左上和右下special_points = [[3400, 1655], # Example points, replace with your own[4550, 2350]]# 讀取圖片image_path = r'D:\data\拉伸原始圖像.jpg'image = cv2.imread(image_path)# 獲取圖像寬度和高度height, width = image.shape[:2]# 定義頂點坐標,你會保留這四個點之內的圖像,以便進行后續步驟trapezoid_vertices = np.array([[0 + generate_random_cut(width), 0 + generate_random_cut(height)],[width + generate_random_cut((width)), 0 + generate_random_cut((height))],[width + generate_random_cut(width), height + generate_random_cut(height)],[0 + generate_random_cut(width), height + generate_random_cut(height)]],dtype=np.float32)# 定義目標圖像的頂點坐標,會將trapezoid_vertices所保留的圖像拉伸,拉伸到target_trapezoid_vertices所對應的范圍內target_trapezoid_vertices = np.array([[0 + generate_random_perspective(width), 0 + generate_random_perspective(height)],[width + generate_random_perspective(width), 0 + generate_random_perspective(height)],[width + generate_random_perspective(width),height + generate_random_perspective(height)],[0 + generate_random_perspective(width), height + generate_random_perspective(height)]],dtype=np.float32)# 計算透視變換矩陣,用于將image進行拉伸perspective_matrix = cv2.getPerspectiveTransform(trapezoid_vertices, target_trapezoid_vertices)# 進行透視變換trapezoid_image = cv2.warpPerspective(image, perspective_matrix, (width, height))# 將坐標框進行處理,利用perspective_matrix得到新的坐標框的左上角和右下角special_points = np.array(special_points, dtype='float32')special_points = np.array([special_points])transformed_special_points = cv2.perspectiveTransform(special_points, perspective_matrix)# 得到新的標注的框的四個頂點的坐標bbox_corners = get_bbox_corners(transformed_special_points)# 如果有超出圖像邊界的點,就計算與圖像邊界的交點,保存到point_list中。point_list = check_intersections(bbox_corners, width, height)# 把留在圖像范圍內的點也加到point_list中point_list.extend(check_points_within_image(bbox_corners, width, height))# 得出新的理想中的標注框的坐標if point_list:rect = get_bounding_box(point_list)else:rect = Noneprint(len(rect))# 就愛那個新的標注框的左上角和右下角繪制出來,以便判斷是否正確for point in rect:x = tuple([int(i) for i in point])cv2.circle(trapezoid_image, x, 15, (256, 256, 256), 10)# 保存變換后的圖像save_path = r'D:\data\拉伸結果.jpg'cv2.imwrite(save_path, trapezoid_image)