Untitled

import os
import grpc
import numpy as np
import random
import tensorflow as tf
import cv2
import glob
import traceback
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:8501",
    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])