Update inference.py

inference.py

@@ -62,10 +62,10 @@ def prediction_mask(prediction_result, map_name, legend, outputPath):
     contour = sorted(contours, key=cv2.contourArea, reverse=True)[0]
     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)
+    mask = cv2.fillPoly(mask, pts=[contour], color=(1)).astype(np.uint8)
 
     # Threshold prediction results and convert to int
-    prediction_result_int = (prediction_result > 0.5).astype(int)
+    prediction_result_int = (prediction_result > 0.5).astype(np.uint8)
 
     # Perform the bitwise operation with the mask also converted to uint8
     masked_img = cv2.bitwise_and(prediction_result_int, mask)