Updated app.py

app.py

@@ -13,11 +13,16 @@ os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'
 os.environ['CUDA_VISIBLE_DEVICES'] = '-1'
 
 # Load labels
-with open('tensorflow/labels.txt', 'r') as f:
+labels_path = os.path.join('tensorflow', 'labels.txt')
+if not os.path.exists(labels_path):
+    raise FileNotFoundError(f"Labels file not found at {labels_path}")
+with open(labels_path, 'r') as f:
     labels = f.read().splitlines()
 
 # Function to load the frozen TensorFlow graph
 def load_frozen_graph(pb_file_path):
+    if not os.path.exists(pb_file_path):
+        raise FileNotFoundError(f"Model file not found at {pb_file_path}")
     with tf.io.gfile.GFile(pb_file_path, 'rb') as f:
         graph_def = tf.compat.v1.GraphDef()
         graph_def.ParseFromString(f.read())
@@ -27,16 +32,19 @@ def load_frozen_graph(pb_file_path):
     return graph
 
 # Load the TensorFlow model
-MODEL_DIR = 'tensorflow/'
+MODEL_DIR = 'tensorflow'
 MODEL_PATH = os.path.join(MODEL_DIR, 'model.pb')
 graph = load_frozen_graph(MODEL_PATH)
 sess = tf.compat.v1.Session(graph=graph)
 
 # Define tensor names based on your model's outputs
-input_tensor = graph.get_tensor_by_name('image_tensor:0')
-detected_boxes = graph.get_tensor_by_name('detected_boxes:0')
-detected_classes = graph.get_tensor_by_name('detected_classes:0')
-detected_scores = graph.get_tensor_by_name('detected_scores:0')
+try:
+    input_tensor = graph.get_tensor_by_name('image_tensor:0')
+    detected_boxes = graph.get_tensor_by_name('detected_boxes:0')
+    detected_classes = graph.get_tensor_by_name('detected_classes:0')
+    detected_scores = graph.get_tensor_by_name('detected_scores:0')
+except KeyError as e:
+    raise KeyError(f"Tensor not found in the graph: {e}")
 
 # Define the target size based on your model's expected input
 TARGET_WIDTH = 320
@@ -53,7 +61,8 @@ def preprocess_image(image):
     image = image.resize((TARGET_WIDTH, TARGET_HEIGHT))
     image_np = np.array(image).astype(np.float32)
     image_np = image_np / 255.0  # Normalize to [0,1]
-    image_np = image_np[:, :, (2, 1, 0)]  # Convert RGB to BGR if required
+    if image_np.shape[2] == 3:
+        image_np = image_np[:, :, (2, 1, 0)]  # Convert RGB to BGR if required
     image_np = np.expand_dims(image_np, axis=0)  # Add batch dimension
     return image_np
 
@@ -128,6 +137,9 @@ def predict(image):
         return annotated_image
     
     except Exception as e:
+        # Log the exception for debugging
+        print(f"Exception during prediction: {e}")
+        
         # Return an error image with the error message
         error_image = Image.new('RGB', (500, 500), color=(255, 0, 0))
         draw = ImageDraw.Draw(error_image)