python - 如何将 TensorFlow 对象检测标签和边界框导出为 CSV 文件？
问题描述
我正在研究对象检测模型以帮助进行检查。一个功能是，如果在输入视频中发现错误，它将发出通知并将数据导出到 CSV。如果可能，我希望它导出错误的标签和边界框坐标，以及时间戳和帧数。
我已经看到了一个关于将边界框导出到 CSV 的答案，但它还没有起作用，并且给了我错误
File "C:\Users\Charles.averill\AppData\Local\Programs\Python\Python36\lib\site-packages\numpy\lib\npyio.py", line 1377, in savetxt
"Expected 1D or 2D array, got %dD array instead" % X.ndim)
ValueError: Expected 1D or 2D array, got 3D array instead
我有它，以便视频注释，一旦完成，它会提示用户是否导出为 CSV，所以我有另一种导出方法。这是我的代码：
def annotate(self):
if("annotated" in self.video_path):
messagebox.showinfo("Error", "You can't annotate an annotated video!")
elif(self.mode == "V" and not self.video_path is None):
fourcc = cv2.VideoWriter_fourcc(*'MP4V')
time = datetime.datetime.now().strftime('%Y-%m-%d %H_%M_%S')
path = 'output/videos/annotated_' + time + '_output.mp4'
out = cv2.VideoWriter(path, fourcc, 20.0, (960, 540))
self.rewind()
NUM_CLASSES = 2
detection_graph = tf.Graph()
with detection_graph.as_default():
od_graph_def = tf.compat.v1.GraphDef()
with tf.io.gfile.GFile(self.model_graph, 'rb') as fid:
serialized_graph = fid.read()
od_graph_def.ParseFromString(serialized_graph)
tf.import_graph_def(od_graph_def, name='')
lmap = label_map_util.load_labelmap(self.label_map)
categories = label_map_util.convert_label_map_to_categories(lmap, max_num_classes=NUM_CLASSES, use_display_name=True)
category_index = label_map_util.create_category_index(categories)
with detection_graph.as_default():
with tf.compat.v1.Session(graph=detection_graph) as sess:
while not self.currentFrame is None:
image_np = self.get_just_frame()
if(image_np is None):
break
image_np_expanded = np.expand_dims(image_np, axis=0)
image_tensor = detection_graph.get_tensor_by_name('image_tensor:0')
self.boxes = detection_graph.get_tensor_by_name('detection_boxes:0')
self.scores = detection_graph.get_tensor_by_name('detection_scores:0')
classes = detection_graph.get_tensor_by_name('detection_classes:0')
num_detections = detection_graph.get_tensor_by_name(
'num_detections:0')
(self.boxes, self.scores, classes, num_detections) = sess.run(
[self.boxes, self.scores, classes, num_detections],
feed_dict={image_tensor: image_np_expanded})
vis_util.visualize_boxes_and_labels_on_image_array(
image_np,
np.squeeze(self.boxes),
np.squeeze(classes).astype(np.int32),
np.squeeze(self.scores),
category_index,
use_normalized_coordinates=True,
line_thickness=2)
# Display output
out.write(image_np)
self.video.release()
out.release()
self.video = None
self.set_video_path(path)
self.video = cv2.VideoCapture(self.video_path)
if(not self.video.isOpened()):
raise ValueError("Unable to open video source", self.video_path)
ret, frame = self.get_frame()
if(ret and not frame is None):
self.photo = PIL.ImageTk.PhotoImage(image = PIL.Image.fromarray(frame))
self.canvas.create_image(0, 0, image = self.photo, anchor = NW)
MsgBox = tk.messagebox.askquestion ('Export to CSV','Do you want to export the video to CSV?',icon = 'warning')
if MsgBox == 'yes':
self.export_CSV()
if(self.video_path is None):
messagebox.showinfo("Error", "No video selected")
def export_CSV(self):
if(not self.boxes is None):
print(self.boxes)
for i, box in enumerate(np.squeeze(self.boxes)):
if(np.squeeze(self.scores)[i] > 0.5):
print("ymin={}, xmin={}, ymax={}, xmax{}".format(box[0]*540,box[1]*960,box[2]*540,box[3]*960))
time = datetime.datetime.now().strftime('%Y-%m-%d %H_%M_%S')
path = 'output/csv/' + time + '_output.csv'
np.savetxt(path, self.boxes, delimiter=',')
else:
messagebox.showinfo("Error", "No boxes, you must\nannotate the video first")
如何将标签与边界框一起导出？
谢谢！
解决方案
您需要将预测的类和类别传递给 writer 方法，也可以使用 csv 库来编写预测。
在文件开头添加：
import csv
我稍微更改了您的代码，但您可以自由更正它。
def annotate(self):
if("annotated" in self.video_path):
messagebox.showinfo("Error", "You can't annotate an annotated video!")
elif(self.mode == "V" and not self.video_path is None):
fourcc = cv2.VideoWriter_fourcc(*'MP4V')
time = datetime.datetime.now().strftime('%Y-%m-%d %H_%M_%S')
path = 'output/videos/annotated_' + time + '_output.mp4'
out = cv2.VideoWriter(path, fourcc, 20.0, (960, 540))
self.rewind()
NUM_CLASSES = 2
detection_graph = tf.Graph()
with detection_graph.as_default():
od_graph_def = tf.compat.v1.GraphDef()
with tf.io.gfile.GFile(self.model_graph, 'rb') as fid:
serialized_graph = fid.read()
od_graph_def.ParseFromString(serialized_graph)
tf.import_graph_def(od_graph_def, name='')
lmap = label_map_util.load_labelmap(self.label_map)
categories = label_map_util.convert_label_map_to_categories(lmap, max_num_classes=NUM_CLASSES, use_display_name=True)
category_index = label_map_util.create_category_index(categories)
with detection_graph.as_default():
with tf.compat.v1.Session(graph=detection_graph) as sess:
while not self.currentFrame is None:
image_np = self.get_just_frame()
if(image_np is None):
break
image_np_expanded = np.expand_dims(image_np, axis=0)
image_tensor = detection_graph.get_tensor_by_name('image_tensor:0')
self.boxes = detection_graph.get_tensor_by_name('detection_boxes:0')
self.scores = detection_graph.get_tensor_by_name('detection_scores:0')
classes = detection_graph.get_tensor_by_name('detection_classes:0')
num_detections = detection_graph.get_tensor_by_name(
'num_detections:0')
(self.boxes, self.scores, classes, num_detections) = sess.run(
[self.boxes, self.scores, classes, num_detections],
feed_dict={image_tensor: image_np_expanded})
vis_util.visualize_boxes_and_labels_on_image_array(
image_np,
np.squeeze(self.boxes),
np.squeeze(classes).astype(np.int32),
np.squeeze(self.scores),
category_index,
use_normalized_coordinates=True,
line_thickness=2)
# Display output
out.write(image_np)
self.video.release()
out.release()
self.video = None
self.set_video_path(path)
self.video = cv2.VideoCapture(self.video_path)
if(not self.video.isOpened()):
raise ValueError("Unable to open video source", self.video_path)
ret, frame = self.get_frame()
if(ret and not frame is None):
self.photo = PIL.ImageTk.PhotoImage(image = PIL.Image.fromarray(frame))
self.canvas.create_image(0, 0, image = self.photo, anchor = NW)
MsgBox = tk.messagebox.askquestion ('Export to CSV','Do you want to export the video to CSV?',icon = 'warning')
if MsgBox == 'yes':
self.export_CSV(self.boxes, self.scores, classes, category_index)
if(self.video_path is None):
messagebox.showinfo("Error", "No video selected")
def export_CSV(self, boxes, scores, classes, category_index):
if (boxes is None):
messagebox.showinfo("Error", "No boxes, you must\nannotate the video first")
return
time = datetime.datetime.now().strftime('%Y-%m-%d %H_%M_%S')
path = 'output/csv/' + time + '_output.csv'
print(boxes)
with open(path, 'wb') as write_file:
writer = csv.writer(write_file)
for box, score, predicted_class in zip(np.squeeze(boxes), np.squeeze(scores), classes):
if(score > 0.5):
print("ymin={}, xmin={}, ymax={}, xmax{}".format(box[0]*540,box[1]*960,box[2]*540,box[3]*960))
writer.writerow([box[0], box[1], box[2], box[3], category_index[predicted_class]['name']])
我不确定 category_index 是否有一个名为“name”的字段，但正如我所说，您可以根据需要随意更改：
writer.writerow([box[0], box[1], box[2], box[3], category_index[predicted_class]['name']])
推荐阅读
- keras - 使用 Keras 分词器来词干
- html - Flexbox 的 row-reverse 属性是不是倒线？
- neo4j - 如何增加 Neo4j 中节点的气泡大小？
- javascript - 如何转置 2x2 数组？
- php - 网格中的 PHP 数据循环
- kubernetes - 如何访问 GKE kubectl 代理仪表板？
- c# - 断言被测系统应该抛出断言异常
- java - 消息在死信队列中结束
- ruby-on-rails - Rails 需要在注册后设置用户的个人资料
- android - 为什么我可以构建和运行我的 Android 应用程序，但无法针对某些 support.v4 和 v7 问题运行单元测试？