from cv2 import cv2
from django.views import View
from django import http
from . import ocrModel
import numpy as np


class InvoiceOCRView(View):
    def post(self, request):
        images = request.FILES.getlist('images', '')
        if images == '':
            return http.JsonResponse({'code': 404, 'results': '请上传图片!!!'})
        invoice_type = int(request.POST.get('invoice_type', 1))
        if invoice_type not in (1, 2):
            return http.JsonResponse({'code': 404, 'results': '输入正确的参数!!!'})
        image_list = []
        for image in images:
            img = np.fromstring(image.read(), np.uint8)
            img = cv2.imdecode(img, cv2.IMREAD_COLOR)
            image_list.append(img)
        ocr = ocrModel
        res = ocr.predict(images=image_list, invoice_type=invoice_type)
        # print(res)
        return http.JsonResponse({'code': 200, 'results': res})