Update inference.py

inference.py

@@ -62,16 +62,14 @@ def prediction_mask(prediction_result, map_name):
     wid, hight = prediction_result.shape[0], prediction_result.shape[1]
     mask = np.zeros([wid, hight])
     mask = cv2.fillPoly(mask, pts=[contour], color=(1,1,1)).astype(int)
-    
-    try:
-        masked_img = cv2.bitwise_and(prediction_result, mask)
-    except cv2.error as e:
-        print("Error with cv2.bitwise_and:")
-        print("prediction_result shape:", prediction_result.shape)
-        print("mask shape:", mask.shape)
-        print("prediction_result dtype:", prediction_result.dtype)
-        print("mask dtype:", mask.dtype)
-        raise e
+
+    # Normalize the float image to be in the range [0, 255] if it's not already
+    prediction_result_normalized = cv2.normalize(prediction_result, None, 0, 255, cv2.NORM_MINMAX)
+    prediction_result_uint8 = prediction_result_normalized.astype(np.uint8)
+
+    # Perform the bitwise operation with the mask also converted to uint8
+    mask_uint8 = mask.astype(np.uint8)
+    masked_img = cv2.bitwise_and(prediction_result_uint8, mask_uint8)
 
     return masked_img