Merge pull request #15 from uug-ai/image-deresolution

Fix NoneType error and add logs

exports/flat/flat_export.py

@@ -62,7 +62,7 @@ def save_frame(self, frame, predicted_frames, cv2, labels_and_boxes, labeled_fra
             f'{self.result_dir_path}/{unix_time}.png',
             frame)
 
-        if labeled_frame.any():
+        if labeled_frame is not None:
             os.makedirs(self.result_labeled_dir_path, exist_ok=True)
 
             cv2.imwrite(

exports/yolov8/yolov8_export.py

@@ -75,7 +75,7 @@ def save_frame(self, frame, predicted_frames, cv2, labels_and_boxes, labeled_fra
             f'{self.image_dir_path}/{unix_time}.png',
             frame)
 
-        if labeled_frame.any():
+        if labeled_frame is not None:
             os.makedirs(self.result_labeled_dir_path, exist_ok=True)
 
             cv2.imwrite(

projects/base_project.py

@@ -44,7 +44,7 @@ def create_proj_save_dir(self):
         _cur_dir = pdirname(pabspath(__file__))
         self.proj_dir = pjoin(_cur_dir, f'../data/{self.name}')
         self.proj_dir = pabspath(self.proj_dir)  # normalise the link
-        print(f'1. Created/Found project folder under {self.proj_dir} path')
+        print(f'3. Created/Found project folder under {self.proj_dir} path')
 
     def connect_models(self):
         """

projects/helmet/helmet_project.py

@@ -27,8 +27,8 @@ def __init__(self):
         super().__init__()
         self._config = self.__read_config__(config_path)
         self.temp_path = self._config.get('temp')
-        self.min_width = int(self._config.get('min_width')) if self._config.get('min_width') else 0
-        self.min_height = int(self._config.get('min_height')) if self._config.get('min_height') else 0
+        self.min_width = int(self._config.get('min_width', '0'))
+        self.min_height = int(self._config.get('min_height', '0'))
         self.models, self.models_allowed_classes = self.connect_models()
         self.mapping = self.class_mapping(self.models)
         self.create_proj_save_dir()
@@ -83,7 +83,7 @@ def class_mapping(self, models):
         result = []
 
         # Iterate through each class index in model_classes[0]
-        for i, class_index in enumerate(model_classes[0]):
+        for class_index in model_classes[0]:
             class_name = model_names[0][class_index]  # Get the class name from the first model
 
             # Create a list to store the mapping for this class
@@ -140,6 +140,6 @@ def connect_models(self):
         if not models:
             raise ModuleNotFoundError('Model not found!')
 
-        print(f'2. Using device: {self.device}')
-        print(f"3. Using {len(models)} models: {[model_name for model_name in self._config.get('models')]}")
+        print(f'1. Using device: {self.device}')
+        print(f"2. Using {len(models)} models: {[model_name for model_name in self._config.get('models')]}")
         return models, models_allowed_classes

services/harvest_service.py

@@ -225,6 +225,7 @@ def predict_frame(self, frame, skip_frames_counter):
             if condition_met:
                 self.predicted_frames = self.export.save_frame(frame, self.predicted_frames, cv2, labels_and_boxes, labeled_frame)
                 skip_frames_counter = self._var.FRAMES_SKIP_AFTER_DETECT
+                print(f'5.3. Done, skipping the next {self._var.FRAMES_SKIP_AFTER_DETECT} frames')
             print(f'Currently in frame: {self.frame_number}')
         self.frame_number += 1
         return skip_frames_counter