Untitled

import os
import grpc
import numpy as np
import random
import tensorflow as tf
import cv2
import glob
import traceback
from PIL import Image, ImageDraw, ImageFont
import matplotlib.pyplot as plt
from tensorflow_serving.apis import predict_pb2
from tensorflow_serving.apis import prediction_service_pb2_grpc

plt.rcParams["figure.figsize"] = [20, 8]

channel = grpc.insecure_channel("192.168.1.198:9502",
    options=[('grpc.max_receive_message_length', 15 * 1024 * 1024)]
)
stub = prediction_service_pb2_grpc.PredictionServiceStub(channel)

def request_tfsv(inputs, outputs, model_name, sig="serving_default", \
                 version_label="stable", version=None):
    grpc_request = predict_pb2.PredictRequest()
    grpc_request.model_spec.name = model_name
    grpc_request.model_spec.signature_name = sig
    if version is None:
        grpc_request.model_spec.version_label = version_label
    else:
        grpc_request.model_spec.version.value = version

    for iname in inputs:
        grpc_request.inputs[iname].CopyFrom(tf.compat.v1.make_tensor_proto(inputs[iname], \
                                                                shape=inputs[iname].shape))
    response = stub.Predict(grpc_request, timeout=60)
    result = []
    for oname in outputs:
        outp = response.outputs[oname]
        shape = [x.size for x in outp.tensor_shape.dim]
        if outp.dtype == 9:
            value = outp.int64_val
        else:
            value = outp.float_val
        result.append(np.reshape(np.array(value), shape))
    return result

def digits_detection(image_bgr):
    img = np.expand_dims(image_bgr[:, :, ::-1], axis=0)
    preds = request_tfsv({"input_tensor": img}, ["detection_boxes", "detection_scores"], "digit_detection")
    boxes = preds[0][0]
    scores = preds[1][0]
    final_boxes = []        
    for sc, (t, l, b, r) in zip(scores, boxes):
        t = int(image_bgr.shape[0] * t)
        b = int(image_bgr.shape[0] * b)
        l = int(image_bgr.shape[1] * l)
        r = int(image_bgr.shape[1] * r)
        if sc > 0.3:
            final_boxes.append([l, t, r, b])
            
    return final_boxes

def image_processing(pil_image, size=(224, 224), grayscale=False, num_channels=3, normalize=True, transpose=True):	
	img_resized = np.array(pil_image.resize(size, resample=Image.BILINEAR))
	if grayscale:
		gray = cv2.cvtColor(img_resized, cv2.COLOR_RGB2GRAY)
		if num_channels == 3:
			img2 = np.zeros_like(img_resized)
			img2[:,:,0] = gray
			img2[:,:,1] = gray
			img2[:,:,2] = gray
		else:
			img2 = np.expand_dims(gray, axis=2)
	else:
		img2 = img_resized
	
	if num_channels == 3:
		if normalize:			
			image = np.divide(np.subtract(img2.astype(np.float32) / 255, np.array([0.485, 0.456, 0.406], dtype=np.float32)), 
                      np.array([0.229, 0.224, 0.225], dtype=np.float32))
		else:
			image = img2 / 255
	else:
		image = img2 / 255

	if transpose:
		image = np.transpose(image, axes=(2, 0, 1))

	return image

def digit_classify(img_rgb):
    img = Image.fromarray(img_rgb)
    image = image_processing(img, (32, 32), True, 1, False)
    # Inference
    image = np.expand_dims(np.array(image), axis=0)
    image = tf.convert_to_tensor(image, dtype=tf.float32)
    preds = request_tfsv({"input": image}, ["output"], "digits")[0][0]
    preds = tf.nn.softmax(preds).numpy()
    return preds.tolist()  

rand = random.choice(files)
image = cv2.imread(rand)
ratio = 64 / image.shape[0]
image = cv2.resize(image, (0, 0), fx=ratio, fy=ratio)
boxes = digits_detection(image)
boxes = sorted(boxes, key=lambda x: x[0])
digits = []
for i, (l, t, r, b) in enumerate(boxes):    
    crop = image.copy()[t:b, l:r, ::-1]    
    pred = digit_classify(crop)
    print(np.argmax(pred), np.max(pred))
    
for i, (l, t, r, b) in enumerate(boxes):
    cv2.rectangle(image, (l, t), (r, b), (0, 255, 0), 1)
    
plt.imshow(image[:, :, ::-1])