前言

最近幫人跑代碼的時候遇到了不少問題，后來發現他提供給我的數據集中標注文件不太準確，部分box框沒有很好地框到物體。所以寫了一個代碼，通過在圖片上繪制標注框來直觀地判斷標注文件是否存在問題，也可以繪制經過網絡預測之后的box框。
參考文檔：python讀取數據集并生成txt文件

一、文件路徑

# -*- coding:utf-8 -*- import os import numpy as np import colorsys import random from PIL import Image,ImageFont, ImageDrawannotation_path = '../model_data/dataset/boats/boats.txt' classes_name = '../model_data/dataset/boats/boat.txt' font_path = '../model_data/simhei.ttf'

annotation_path 是數據集的txt文件索引，classes_name 存放了所有類別名稱，font_path存放的是字體文件（也可以不設置）

二、在圖片上繪制標注框

#獲取所有類別 def get_classes(classes_path):classes = os.path.expanduser(classes_path)with open(classes) as f:class_names = f.readlines()class_names = [c.strip() for c in class_names]return class_namesdef show_ground_truth(annotation_line):print(annotation_line)line = annotation_line.split()#獲取圖片image = Image.open(line[0])#獲取GT_box框boxes = np.array([np.array(list(map(int, box.split(',')))) for box in line[1:]])#獲取類別數class_names = get_classes(classes_name)# 生成繪制邊框的顏色。# h(色調）：x/len(self.class_names) s(飽和度）：1.0 v(明亮）：1.0# 對于所有種類的目標，確定每一種目標框的繪制顏色，即：將(x/80, 1.0, 1.0)的顏色轉換為RGB格式，并隨機調整顏色以便于肉眼識別，hsv_tuples = [(x / len(class_names), 1., 1.) for x in range(len(class_names))]colors = list(map(lambda x: colorsys.hsv_to_rgb(*x), hsv_tuples)) # hsv轉換為rgb# hsv取值范圍在[0,1]，而RBG取值范圍在[0,255]，所以乘上255colors = list(map(lambda x: (int(x[0] * 255), int(x[1] * 255), int(x[2] * 255)),colors))# 定義字體font = ImageFont.truetype(font=font_path,size=np.floor(3e-2 * np.shape(image)[1] + 0.5).astype('int32'))# 設置目標框線條的寬度model_image_size=[416,416,3]thickness = (np.shape(image)[0] + np.shape(image)[1]) // model_image_size[0]for i in range(len(boxes)):# 創建一個可以在給定圖像上繪圖的對象draw = ImageDraw.Draw(image)#獲取目標框left, top, right, bottom, predicted_class = boxes[i]predicted_class = class_names[predicted_class]label = '{}'.format(predicted_class)top = top - 5left = left - 5bottom = bottom + 5right = right + 5# 防止目標框溢出top = max(0, np.floor(top + 0.5).astype('int32'))left = max(0, np.floor(left + 0.5).astype('int32'))bottom = min(np.shape(image)[0], np.floor(bottom + 0.5).astype('int32'))right = min(np.shape(image)[1], np.floor(right + 0.5).astype('int32'))#定義標簽文字label_size = draw.textsize(label, font)label = label.encode('utf-8')# 確定標簽（label）起始點位置：標簽的左、下if top - label_size[1] >= 0:text_origin = np.array([left, top - label_size[1]])else:text_origin = np.array([left, top + 1])#畫目標框，線條寬度為thickness(這里畫了很多個框重疊起來)for i in range(thickness):draw.rectangle([left + i, top + i, right - i, bottom - i],outline=colors[class_names.index(predicted_class)])#畫標簽框# draw.rectangle(# [tuple(text_origin), tuple(text_origin + label_size)],# fill=colors[class_names.index(predicted_class)])# 填寫標簽內容# draw.text(text_origin, str(label, 'UTF-8'), fill=(0, 0, 0), font=font)del draw#展示圖片image.show()

1、按空格將annotation_line分隔開，獲取圖片的路徑和所有標注框的坐標和類別信息。
2、定義繪制邊框的顏色，將hsv格式轉換成RGB格式。定義字體樣式。設置目標框線條的寬度（為了保證線條的寬度在不同的圖片上看起來都比較合適，這里將寬度設置為圖片的長寬加起來除以416）。
3、利用for循環在圖片上依次畫框。為防止目標框溢出，目標框的上、左兩個坐標小數點后一位四舍五入；目標框的下、右兩個坐標小數點后一位四舍五入，與圖片的尺寸相比，取最小值。
4、最后畫目標框（這里實際上畫了很多框聚在一起，所以看起來才有寬度）和標簽框（如果目標太多顯得擁擠可以將代碼注釋掉），顯示圖片。

三、主函數

if __name__ == '__main__':file = open(annotation_path)lines = file.read().splitlines()#隨機讀取一張圖片#show_ground_truth(random.choice(lines))#讀取所有圖片for line in lines:show_ground_truth(random.choice(lines))os.system("pause");

既可以隨機讀取一張圖片，也可以遍歷所有圖片。

四、總結

如果你在處理過程中遇到了任何問題，都可以在下方進行留言和評論。

總結

以上是生活随笔為你收集整理的python在图片上绘制标注框的全部內容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網站內容還不錯，歡迎將生活随笔推薦給好友。