Update app.py

app.py

@@ -98,9 +98,15 @@ def predict(radio, dict, word_mask, prompt=""):
         init_image = dict['image'].convert('RGB').resize((imgRes, imgRes))
         filename = f"{uuid.uuid4()}.png"
         plt.imsave(filename,torch.sigmoid(preds[0][0]))
-        img2 = cv2.imread(filename)
+        
+        img = cv2.imread(filename, cv2.IMREAD_GRAYSCALE)
+        img = Image.fromarray(img)
+
+        #img2 = cv2.imread(filename)
+        
         #if ret == True: 
-        gray_image = cv2.cvtColor(img2, cv2.COLOR_BGR2GRAY)
+        gray_image = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
+        
         (thresh, bw_image) = cv2.threshold(gray_image, 100, 255, cv2.THRESH_BINARY)
         cv2.cvtColor(bw_image, cv2.COLOR_BGR2RGB)
         mask = Image.fromarray(np.uint8(bw_image)).convert('RGB')
@@ -128,10 +134,13 @@ def predict(radio, dict, word_mask, prompt=""):
         init_image = dict['image'].convert('RGB').resize((imgRes, imgRes))
         filename = f"{uuid.uuid4()}.png"
         plt.imsave(filename,torch.sigmoid(preds[0][0]))
-        img2 = cv2.imread(filename)
+        #img2 = cv2.imread(filename)
+        
+        img = cv2.imread(filename, cv2.IMREAD_GRAYSCALE)
+        img = Image.fromarray(img)
         
         #if ret == True: 
-        gray_image = cv2.cvtColor(img2, cv2.COLOR_BGR2GRAY)
+        gray_image = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
         (thresh, bw_image) = cv2.threshold(gray_image, 100, 255, cv2.THRESH_BINARY)
         cv2.cvtColor(bw_image, cv2.COLOR_BGR2RGB)
         mask = Image.fromarray(np.uint8(bw_image)).convert('RGB')