Update main.py

main.py

@@ -1,9 +1,16 @@
+import tensorflow as tf
+
+# Detect and init the GPU
+physical_devices = tf.config.list_physical_devices('GPU')
+if len(physical_devices) > 0:
+    tf.config.set_visible_devices(physical_devices[0], 'GPU')
+    tf.config.experimental.set_memory_growth(physical_devices[0], True)
+
 import io
 from flask import Flask, request, jsonify
 import base64
 import pytesseract
 import numpy as np
-import tensorflow as tf
 from pickle import load
 from PIL import Image
 from keras.applications.xception import Xception # to get pre-trained model Xception
@@ -14,17 +21,6 @@ app = Flask(__name__)
 
 MAX_LENGTH = 38
 
-# Set up GPU memory growth
-physical_devices = tf.config.list_physical_devices('GPU')
-if physical_devices:
-    try:
-        # Allow memory growth for all GPUs
-        for gpu in physical_devices:
-            tf.config.experimental.set_memory_growth(gpu, True)
-        print("GPU(s) memory growth set to True")
-    except RuntimeError as e:
-        print(e)
-
 def format_tesseract_output(output_text):
     formatted_text = ""
     lines = output_text.strip().split("\n")
@@ -39,8 +35,6 @@ def extract_features(image_data, model):
     try:
         image = Image.open(io.BytesIO(image_data))
     except Exception as e:
-        print("ERROR: Can't open image! Ensure that image data is correct and in the expected format")
-        print(str(e))
         return None
 
     image = image.resize((299,299))