new update

Segmentation.py

@@ -10,6 +10,7 @@ def black_mask(results):
     if not os.path.exists(runs_folder):
         # Nếu chưa tồn tại, tạo thư mục "mask"
         os.makedirs(runs_folder)
+    bl_mask=[]
     for mask_index in range(len(results[0].masks)):
         # Lấy mặt nạ
             mask_raw = results[0].masks[mask_index].cpu().data.numpy().transpose(1, 2, 0)
@@ -34,9 +35,11 @@ def black_mask(results):
             black_color = [0, 0, 0]
             white_color = [255, 255, 255]
             result_image = np.where(np.all(masked == black_color, axis=-1, keepdims=True), masked, white_color)
-            name = "mask/"+str(mask_index) + '.png'
-            # Show the masked part of the image
-            cv2.imwrite(name,result_image)
+            # name = "mask/"+str(mask_index) + '.png'
+            # # Show the masked part of the image
+            # cv2.imwrite(name,result_image)
+            bl_mask.append(result_image)
+    return bl_mask
 def segmentation_doc(image):
     model = YOLO('Model/Seg65ep.pt')
     runs_folder = "runs"
@@ -52,7 +55,7 @@ def segmentation_doc(image):
     padding_size = 50
     image = cv2.copyMakeBorder(image, padding_size, padding_size, padding_size, padding_size, cv2.BORDER_CONSTANT, value=color)
     results = model(image,save=True,retina_masks = True, conf=0.7)
-    black_mask(results)
-    return image
+    bl_mask=black_mask(results)
+    return image,bl_mask
 # image=cv2.imread("test1.jpg")
 # segmentation_doc(image)

app.py

@@ -8,13 +8,13 @@ from Rotate.rotate_function import *
 from Rotate import rotate_function
 # from Restomer import *
 def process_image(input_image):
-    original_img=segmentation_doc(input_image)
+    original_img,masks=segmentation_doc(input_image)
     output_images=[]
-    mask_path="mask"
-    masks = [f for f in os.listdir(mask_path) if os.path.isfile(os.path.join(mask_path, f))]
-    for mask_name in masks:
-        mask=os.path.join(mask_path, mask_name)
-        mask_img=cv2.imread(mask)
+    # mask_path="mask"
+    # masks = [f for f in os.listdir(mask_path) if os.path.isfile(os.path.join(mask_path, f))]
+    for mask_img in masks:
+        # mask=os.path.join(mask_path, mask_name)
+        # mask_img=cv2.imread(mask)
         extract_document=extract(original_img,mask_img)
         # result=rotate_function.rotate(extract_document)
         # result=run_script(result, 720, 32)