【TensorFlow/Python】おすすめアノテーションツール VoTT で自動的にアノテーション作業とタグづけを行うPythonコードとVoTTの設定方法。〜Object Detect 用訓練data自動作成によるActive Learning〜

import sys
import json
import hashlib
import urllib
from pathlib import Path
from PIL import Image # pillow 7.1.2
import numpy as np
import os
import six.moves.urllib as urllib
import sys
import tarfile
import tensorflow as tf
import zipfile

from collections import defaultdict
from io import StringIO
from matplotlib import pyplot as plt
#from PIL import Image
sys.path.append("..")

from models.research.object_detection.utils import label_map_util
from models.research.object_detection.utils import visualization_utils as vis_util
import cv2
import glob
##########モジュールインポート##########

##########任意モデルのパス#############
MODEL_NAME = '/models/research/object_detection/ssd_mobilenet_v1_coco_2018_01_28/'
PATH_TO_CKPT = os.path.join(MODEL_NAME, 'frozen_inference_graph.pb')
TFRec_Dir = '/models/research/object_detection/data/'
PATH_TO_LABELS = os.path.join(TFRec_Dir, 'mscoco_label_map.pbtxt')

NUM_CLASSES = 90
####################################

##########tensorflowの動作コード######
detection_graph = tf.Graph()
with detection_graph.as_default():
    od_graph_def = tf.compat.v1.GraphDef()
    with tf.io.gfile.GFile(PATH_TO_CKPT, 'rb') as fid:
        serialized_graph = fid.read()
        od_graph_def.ParseFromString(serialized_graph)
        tf.import_graph_def(od_graph_def, name='')

label_map = label_map_util.load_labelmap(PATH_TO_LABELS)
categories = label_map_util.convert_label_map_to_categories(label_map, max_num_classes=NUM_CLASSES, use_display_name=True)
category_index = label_map_util.create_category_index(categories)


########################################
imgs = glob.glob('/vott/test_image/*.jpg')
##########インプット画像path###############

def make_point(img_path):
    height = 840
    width  = 1240
    image_np = cv2.imread(img_path)
    original_image = image_np.copy()
    image_width, image_height = original_image.shape[1] , original_image.shape[0]

    with detection_graph.as_default():

        with tf.compat.v1.Session(graph=detection_graph) as sess:

            image_np_expanded = np.expand_dims(image_np, axis=0)
            image_tensor = detection_graph.get_tensor_by_name('image_tensor:0')

            boxes = detection_graph.get_tensor_by_name('detection_boxes:0')
            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')

            (boxes, scores, classes, num_detections) = sess.run(
                [boxes, scores, classes, num_detections],
                feed_dict={image_tensor: image_np_expanded})
                
            vis_util.visualize_boxes_and_labels_on_image_array(
                image_np,
                np.squeeze(boxes),
                np.squeeze(classes).astype(np.int32),
                np.squeeze(scores),
                category_index,
                use_normalized_coordinates=True,
                line_thickness=8)
            
            real_boxes = np.squeeze(boxes)
            real_classes = np.squeeze(classes).astype(np.int32)
            real_scores = np.squeeze(scores)
            
            global y_min, x_min, y_max, x_max
            count = 0
            vott_list = []
            for item in real_boxes:
                if real_classes[count] == 18:

                    vott_list.append(real_scores[count])

                else:
                    pass
                count += 1
            print(vott_list.index(max(vott_list)))
            top_number = vott_list.index(max(vott_list))
            vott_target = real_boxes[top_number]
            ymin, xmin, ymax, xmax = vott_target[0] ,vott_target[1] ,vott_target[2] ,vott_target[3]
            (x_min, x_max, y_min, y_max) = (xmin * image_width, xmax * image_width, ymin * image_height, ymax * image_height)
            
            return(x_min, x_max, y_min, y_max)
            ##出力は物体領域4点を出力

####################################

#########Vottの作成するJsonの作成######
for img in imgs:
    make_point(img)
    imgdir = Path(sys.argv[1]).resolve()
    projfile = imgdir.with_suffix(".json")

    if projfile.exists():
        # 既存のプロジェクトを読み込む
        f = open(projfile, "r+")
        data = json.load(f)
    else:
        # プロジェクトを新規作成する
        f = open(projfile, "w")
        data = {
            "frames": {},
            "framerate": "1",
            "inputTags": "mrfk", # タグリスト（カンマ区切り）
            "suggestionType": "track",
            "scd": False,
            "visitedFrames": [],
            "tag_colors": ["#0cc7ff"] # 領域の色（任意）
        }

    with f:
        for imgfile in imgdir.glob("*.*"):
            if imgfile.suffix.lower() in [".jpg", ".png"]:
                if imgfile.name in data["frames"]:
                    # エントリがある場合はスキップ
                    continue
                else:
                    # エントリがない場合、新規作成
                    frame = []
                    data["frames"][imgfile.name] = frame

                # 画像のサイズを取得
                img = Image.open(imgfile)
                w, h = img.size
                img.close()

############## 自分で検出した領域を列挙する##############
                points = [
                    {"x": x_min, "y": y_min},
                    {"x": x_min, "y": y_max},
                    {"x": x_max, "y": y_min},
                    {"x": x_max, "y": y_max}
                ]
####################################################

                box = {  # 外接長方形
                    "x1": min(p["x"] for p in points),
                    "y1": min(p["y"] for p in points),
                    "x2": max(p["x"] for p in points),
                    "y2": max(p["y"] for p in points)
                }
                region = box.copy()
                region.update({
                    "width": w,
                    "height": h,
                    "box": box,
                    "points": points,
                    "type": "rect",
                    "tags": ["mrfk"], # 領域に付与するタグリスト
                })
                frame.append(region)

        # プロジェクトを保存
        f.seek(0)
        json.dump(data, f)

python vott_auto.py [指定のimgフォルダpath]