Merge branch 'master' of https://github.com/sof-danny/spannotation

README.md

@@ -40,7 +40,7 @@ or
 
 If you want to install a specific version:
 
-``` pip install spannotation==0.1.9 ```
+``` pip install spannotation==0.1.11 ```
 
 
 ## Usage on Python code Editor
@@ -60,7 +60,7 @@ or
 
 If you want to install a specific version:
 
-``` pip install spannotation==0.1.9 ```
+``` pip install spannotation==0.1.11 ```
 
 
 

Spannotation/cli.py

@@ -23,3 +23,7 @@ def main():
 
 if __name__ == "__main__":
     main()
+
+
+
+

build/lib/Spannotation/mask_generator.py

@@ -3,15 +3,18 @@
 import os
 import glob
 
+
 class MaskGenerator:
     def __init__(self):
         self.points = []
         self.image = None
         self.mask_generated = False
+        self.window_name = 'Image Mask Generator'  
+
 
     def click_event(self, event, x, y, flags, param):
         if event == cv2.EVENT_LBUTTONDOWN:
-            if len(self.points) < 3:
+            if len(self.points) < 3:  
                 self.points.append((x, y))
                 cv2.circle(self.image, (x, y), 5, (0, 255, 0), -1)
                 cv2.imshow('image', self.image)
@@ -20,35 +23,41 @@ def click_event(self, event, x, y, flags, param):
                 mask = self.generate_mask(self.points)
                 cv2.imshow('mask', mask)
                 self.mask_generated = True
+                self.save_mask(mask, param['save_path'], param['image_path'])
+                cv2.waitKey(500)  
+                cv2.destroyAllWindows()
+                cv2.waitKey(1)  
 
-    @staticmethod
-    def generate_mask(image, points):
-        mask = np.zeros(image.shape[:2], dtype=np.uint8)
+    def generate_mask(self, points):
+        mask = np.zeros(self.image.shape[:2], dtype=np.uint8)
         cv2.fillPoly(mask, [np.array(points)], 255)
         return mask
 
+    def save_mask(self, mask, save_path, image_path):
+        base_name = os.path.basename(image_path)
+        save_file = os.path.join(save_path, base_name)
+        cv2.imwrite(save_file, mask)
+        print(f"Mask saved to {save_file}")
+
+
     def process_image(self, image_path, save_path):
         self.image = cv2.imread(image_path)
-        if self.image is None:
-            raise ValueError("Could not read the image.")
         self.points = []
         self.mask_generated = False
 
         cv2.imshow('image', self.image)
-        cv2.setMouseCallback('image', self.click_event)
-        cv2.waitKey(0)
+        cv2.setMouseCallback('image', self.click_event, param={'image_path': image_path, 'save_path': save_path})
+
+        while not self.mask_generated:
+            if cv2.waitKey(1) & 0xFF == ord('q'): 
+                break
 
-        if self.mask_generated:
-            mask = self.generate_mask(self.image, self.points)
-            base_name = os.path.basename(image_path)
-            cv2.imwrite(os.path.join(save_path, base_name), mask)
+        cv2.destroyAllWindows()
+        cv2.waitKey(1)  
 
     def process_folder(self, folder_path, save_path):
         if not os.path.exists(save_path):
             os.makedirs(save_path)
 
         for image_file in glob.glob(os.path.join(folder_path, '*')):
-            try:
-                self.process_image(image_file, save_path)
-            except ValueError as e:
-                print(f"Error processing {image_file}: {e}")
+            self.process_image(image_file, save_path)