Update app.py

app.py

@@ -165,8 +165,14 @@ def inference(image):
     scratch_masks_arr = np.array(scratch_masks)
     damage_masks_arr = np.array(damage_masks)
     # Resize the scratch and damage masks to match the size of the original image
-    scratch_mask_resized = cv2.resize(scratch_masks_arr[0].astype(np.uint8), (img.shape[1], img.shape[0]))
-    damage_mask_resized = cv2.resize(damage_masks_arr[0].astype(np.uint8), (img.shape[1], img.shape[0]))
+    if len(scratch_masks_arr) > 0:
+      scratch_mask_resized = cv2.resize(scratch_masks_arr[0].astype(np.uint8), (img.shape[1], img.shape[0]))
+    else:
+      scratch_mask_resized = np.zeros((img.shape[0], img.shape[1]), dtype=np.uint8)
+    if len(damage_masks_arr) > 0:
+      damage_mask_resized = cv2.resize(damage_masks_arr[0].astype(np.uint8), (img.shape[1], img.shape[0]))
+    else:
+      damage_mask_resized = np.zeros((img.shape[0], img.shape[1]), dtype=np.uint8)
     # Merge the scratch and damage masks into a single binary mask
     merged_mask = np.zeros_like(scratch_mask_resized)
     merged_mask[(scratch_mask_resized> 0) | (damage_mask_resized > 0)] = 255
@@ -243,6 +249,7 @@ def inference(image):
 with gr.Blocks() as demo:
     with gr.Row():
         with gr.Column():
+            gr.HTML("<h1 style='text-align: center;'>Damage Detection Dashboard</h1>")
             gr.Markdown("## Inputs")
             image = gr.Image(type="pil",label="Input")
             submit_button = gr.Button(value="Submit", label="Submit")