import cv2
import numpy as np
import argparse
import imutils
from PIL import Image
import pytesseract
from transform import four_point_transform
import time

config = ('-l eng --oem 1 --psm 6')

ball_color = 'yellow'

color_dist = {'red': {'Lower': np.array([0, 60, 60]), 'Upper': np.array([6, 255, 255])},
              'blue': {'Lower': np.array([100, 80, 46]), 'Upper': np.array([124, 255, 255])},
              'green': {'Lower': np.array([35, 43, 35]), 'Upper': np.array([90, 255, 255])},
              'yellow': {'Lower': np.array([26, 43, 46]), 'Upper': np.array([34, 255, 255])}

              }

cap = cv2.VideoCapture(0)
cv2.namedWindow('camera', cv2.WINDOW_AUTOSIZE)
index = 0
x1 = 0
x2 = 0
y2 = 0
y1 = 0
while cap.isOpened():
    ret, frame = cap.read()
    img_org2 = frame.copy()

    input = cv2.waitKey(1) & 0xFF
    if ret:
        if frame is not None:
            gs_frame = cv2.GaussianBlur(frame, (5, 5), 0)  # 高斯模糊
            hsv = cv2.cvtColor(gs_frame, cv2.COLOR_BGR2HSV)  # 轉(zhuǎn)化成HSV圖像
            erode_hsv = cv2.erode(hsv, None, iterations=2)  # 腐蝕 粗的變細(xì)
            inRange_hsv = cv2.inRange(erode_hsv, color_dist[ball_color]['Lower'], color_dist[ball_color]['Upper'])
            cnts = cv2.findContours(inRange_hsv.copy(), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)[-2]

            cnts = sorted(cnts, key=cv2.contourArea, reverse=True)
            for c in cnts:
                rect = cv2.minAreaRect(c)
                box = cv2.boxPoints(rect)
                box = np.int0(box)
                # x, y, w, h = cv2.boundingRect(c)
                ys = [box[0, 1], box[1, 1], box[2, 1], box[3, 1]]
                xs = [box[0, 0], box[1, 0], box[2, 0], box[3, 0]]
                ys_sorted_index = np.argsort(ys)
                xs_sorted_index = np.argsort(xs)

                x1 = box[xs_sorted_index[0], 0]
                x2 = box[xs_sorted_index[3], 0]

                y1 = box[ys_sorted_index[0], 1]
                y2 = box[ys_sorted_index[3], 1]

                cv2.drawContours(frame, [box], -1, (255, 255, 0), 2)
                cv2.imshow("camera", frame)

                break

            # cv2.waitKey(1)
            if input == ord('p'):
                time1 = time.time()
                cv2.imwrite("%s/%d.png" % ('./', index),
                            cv2.resize(frame, (224, 224), interpolation=cv2.INTER_AREA))

                pts1 = np.float32([[x1,y1],[x2,y1],[x1,y2],[x2,y2]])
                pts2 = np.float32([[0,38],[0,0],[139,38],[139,0]])

                M = cv2.getPerspectiveTransform(pts1,pts2)
                dst = cv2.warpPerspective(img_org2,M,(139,38))
                img_plate = img_org2[y1:y2, x1:(x2-80)]

                print("%s: %d 張圖片" % ('./', index))
                cv2.imwrite("%s/%d_process.png" % ('./', index),
                            cv2.resize(img_plate, (139, 38), interpolation=cv2.INTER_AREA))

                text = pytesseract.image_to_string(img_plate, config=config)
                filter(str.isalnum,text)
                print(text)
                print('The code run{:f}s'.format(time.time()-time1))
                cv2.putText(frame, text, (x1, y1 + (y2 - y1) + 12), cv2.FONT_HERSHEY_COMPLEX, 0.3, (0, 0, 255), 1,
                            cv2.LINE_AA)
                cv2.imshow("camera", frame)


                index += 1
            if input == ord('q'):
                break

        else:
            print("無畫面")
    else:
        print("無法讀取攝像頭！")

cap.release()
cv2.waitKey(0)
cv2.destroyAllWindows()