#!/usr/bin/python

from detectron2.config import get_cfg
from detectron2.engine import DefaultTrainer
from detectron2.engine import DefaultPredictor
from pathlib import Path
import os
import cgi
import base64
import cv2
import warnings
import numpy
import time

# 検出された物体の情報を格納するクラス
class DetectedObject:

	left = 0
	top = 0

	width = 0
	height = 0

	score = 0
	obj_class = 0

	mask = None

	# 画像全体のマスクデータから検出領域のマスクデータを作成
	def setPredMask(self, pred_mask):

		min_x = 9999
		min_y = 9999
		max_x = -1
		max_y = -1

		for y in range(len(pred_mask)):
			for x in range(len(pred_mask[y])):
				if pred_mask[y][x]:

					if x < min_x:
						min_x = x

					if x > max_x:
						max_x = x

					if y < min_y:
						min_y = y

					if y > max_y:
						max_y = y

		self.left = min_x
		self.top = min_y
		self.width = max_x - min_x + 1
		self.height = max_y - min_y + 1

		mask_flag_length = self.width * self.height
		mask_flag_data_length = mask_flag_length + (-mask_flag_length % 8)

		mask_flags = [0] * mask_flag_data_length

		# 画像全体のマスクデータを1次元配列化
		for y in range(self.height):
			for x in range(self.width):
				if pred_mask[y + self.top][x + self.left]:
					mask_flags[x + y * self.width] = 1

		mask_data_items = mask_flag_data_length // 8
		mask_data = bytearray()

		# 8ビット（ピクセル）毎にまとめてバイナリデータに変換
		for i in range(mask_data_items):

			v = 0

			for j in range(8):

				index = j + i * 8

				if (index < mask_flag_data_length) and (mask_flags[index] == 1):
					v += 1 << j

			mask_data.append(v)

		self.mask = base64.b64encode(mask_data).decode()

	def getLeft(self):
		return self.left

	def getTop(self):
		return self.top

	def getWidth(self):
		return self.width

	def getHeight(self):
		return self.height

	def getMask(self):
		return self.mask

	def getClass(self):
		return self.obj_class

	def setClass(self, obj_class):
		self.obj_class = obj_class

	def getScore(self):
		return self.score

	def setScore(self, score):
		self.score = score

def createDetectedObjectList(instances):

	instance_num = len(instances.scores)

	if instance_num == 0:
		return None

	object_list = []

	scores = instances.scores
	obj_classes = instances.pred_classes

	pred_masks = instances.pred_masks
	pred_masks = numpy.asarray(pred_masks)

	for i in range(instance_num):

		obj = DetectedObject()

		obj.setScore(scores[i].item())
		obj.setClass(obj_classes[i].item())

		obj.setPredMask(pred_masks[i])

		object_list.append(obj)

	return object_list

def detectImage(predictor, img):

	output = predictor(img)

	instances = output["instances"].to("cpu")

	return createDetectedObjectList(instances)

def createDetectedObjestsJson(predictor, img):

	# 物体検出処理で物体リストを作成
	detected_objects_list = detectImage(predictor, img)

	if detected_objects_list == None:
		return "\"object_list\": null"

	obj_json = "\"object_list\":["

	# 各物体の情報をJSONデータにする
	for i in range(len(detected_objects_list)):

		obj = detected_objects_list[i]

		if i > 0:
			obj_json = obj_json + ", "

		obj_json = obj_json + "{\"object_type\": " + str(obj.getClass()) + ", ";
		obj_json = obj_json + "\"score\": " + str(obj.getScore()) + ", ";

		obj_json = obj_json + "\"left\": " + str(obj.getLeft()) + ", "
		obj_json = obj_json + "\"top\": " + str(obj.getTop()) + ", "
		obj_json = obj_json + "\"width\": " + str(obj.getWidth()) + ", "
		obj_json = obj_json + "\"height\": " + str(obj.getHeight()) + ", "

		obj_json = obj_json + "\"mask\": \"" + obj.getMask() + "\""

		obj_json = obj_json + "}"

	obj_json = obj_json + "]"

	return obj_json

def cgi_main():

	warnings.simplefilter('ignore')
	params = cgi.FieldStorage()

	# Detectron2処理中ファイルのパス
	processing_file_path = "./detectron_processing"

	# Detectron2処理中ファイルが存在していたら中断
	if os.path.isfile(processing_file_path):

		# Detectron2処理中ファイルが作成されてからの経過時間を算出
		processing_start_sec = os.path.getctime(processing_file_path)
		processing_sec = time.time() - processing_start_sec

		result_status = "busy"

		# 処理が長引いていたら、クライアント側プログラムからの処理中止を促す
		if processing_sec > 250:
			result_status = "error"

		print("Content-Type: text/plain;")
		print("")

		print("{\"result\": \"ng\", \"status\": \"" + result_status + "\"}")

		return

	# Detectron2処理中ファイルを作成
	Path(processing_file_path).touch()

	image_data = params.getfirst("image_data")

	img_buf = base64.b64decode(image_data)
	img_bin = numpy.frombuffer(img_buf, dtype=numpy.uint8)

	img = cv2.imdecode(img_bin, cv2.IMREAD_COLOR)

	cfg = get_cfg()
	cfg.MODEL.DEVICE='cpu'
	cfg.merge_from_file("./detectron2/configs/COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml")
	cfg.DATALOADER.NUM_WORKERS = 2
	cfg.MODEL.WEIGHTS = "detectron2://COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x/137849600/model_final_f10217.pkl"

	cfg.MODEL.WEIGHTS = os.path.join("buttocks/buttocks.pth")
	cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.5
	cfg.MODEL.ROI_HEADS.NUM_CLASSES = 2

	predictor = DefaultPredictor(cfg)

	object_list_json = createDetectedObjestsJson(predictor,img)

	print("Content-Type: text/plain;")
	print("")

	print("{\"result\": \"ok\", " + object_list_json + "}")

	# Detectron2処理中ファイルを削除
	if os.path.isfile(processing_file_path):
		os.remove(processing_file_path)

cgi_main()