How to export TensorFlow object detection labels and bounding boxes as CSV file?

Question

I'm working on an object detection model to aid in inspection. One feature is that if it finds an error in the input video, it will raise a notification and export the data to CSV. I'd like it to export the label and bounding box coordinates of the error, along with the timestamp and frame count if possible.

I've seen one answer about exporting bounding boxes to CSV but it hasn't worked yet, and gives me the error

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

I have it so that the video is annotating, and once it's finished, it prompts the user to export to CSV or not, so I have another method for the exporting. Here's my code:

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")

How could I export the labels along with the bounding boxes?

Thank you!

Answer 1

You need to pass predicted classes and categories to the writer method also you could use csv library to write the predictions.

add at the beginning of the file:

import csv

I changed your code a little but you're free to correct it.

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']])

I'm not sure about category_index that it has a field called "name" but as I said you are free to change to you're needs:

writer.writerow([box[0], box[1], box[2], box[3], category_index[predicted_class]['name']])