Merge pull request #8 from uug-ai/enhancement

update dockerfile

Dockerfile

@@ -43,8 +43,19 @@ COPY . .
 ENV MEDIA_SAVEPATH "/ml/data/input/input_video.mp4"
 
 # Model parameters
-ENV MODEL_NAME: "helmet_dectector_1k_16b_150e.pt"
+ENV MODEL_NAME: "yolov8n.pt.pt"
 ENV MODEL_NAME_2: "helmet_dectector_1k_16b_150e.pt"
+ENV MODEL_ALLOWED_CLASSES="0"
+ENV MODEL_2_ALLOWED_CLASSES="0"
+
+# Dataset parameters
+ENV DATASET_FORMAT="base"
+ENV DATASET_VERSION="1"
+ENV DATASET_UPLOAD="True"
+
+# Forwarding
+ENV FORWARDING_MEDIA="True"
+ENV REMOVE_AFTER_PROCESSED="True"
 
 # Queue parameters
 ENV QUEUE_NAME ""
@@ -59,12 +70,20 @@ ENV STORAGE_URI ""
 ENV STORAGE_ACCESS_KEY ""
 ENV STORAGE_SECRET_KEY ""
 
+#Integration parameters
+ENV INTEGRATION_NAME=""
+
 # Roboflow parameters
-ENV RBF_UPLOAD: ""
 ENV RBF_API_KEY: ""
 ENV RBF_WORKSPACE: ""
 ENV RBF_PROJECT: ""
 
+#S3 parameters
+ENV S3_ENDPOINT=""
+ENV S3_ACCESS_KEY=""
+ENV S3_SECRET_KEY=""
+ENV S3_BUCKET=""
+
 # Feature parameters
 ENV PLOT "False"
 

exports/base_export.py

@@ -10,7 +10,15 @@
 
 
 class BaseExport(IBaseExport):
+    """
+    Base Export class that implements functions for
+    initializing and saving frame under specific format.
+    """
+
     def __init__(self, proj_dir_name):
+        """
+        Constructor.
+        """
         self._var = VariableClass()
         _cur_dir = pdirname(pabspath(__file__))
         self.proj_dir = pjoin(_cur_dir, f'../data/{proj_dir_name}')
@@ -19,10 +27,10 @@ def __init__(self, proj_dir_name):
 
     def initialize_save_dir(self):
         """
-        See ibase_project.py
+        See ibase_export.py
 
         Returns:
-            None
+            success True or False
         """
         self.result_dir_path = pjoin(self.proj_dir, f'{self._var.DATASET_FORMAT}-v{self._var.DATASET_VERSION}')
         os.makedirs(self.result_dir_path, exist_ok=True)
@@ -35,6 +43,12 @@ def initialize_save_dir(self):
             return False
 
     def save_frame(self, frame, predicted_frames, cv2, labels_and_boxes):
+        """
+        See ibase_export.py
+
+        Returns:
+            Predicted frame counter.
+        """
         print(f'5.1. Condition met, processing valid frame: {predicted_frames}')
         # Save original frame
         unix_time = int(time.time())

exports/export_factory.py

@@ -7,6 +7,7 @@ class ExportFactory:
     """
     Export Factory initializes specific export types.
     """
+
     def __init__(self):
         self._var = VariableClass()
         self.save_format = self._var.DATASET_FORMAT

exports/ibase_export.py

@@ -2,11 +2,27 @@
 
 
 class IBaseExport(ABC):
+    """
+    Interface for Base Export.
+    """
 
     @abstractmethod
     def initialize_save_dir(self):
+        """
+        Initializes save directory for Base export format
+        """
         pass
 
     @abstractmethod
     def save_frame(self, frame, predicted_frames, cv2, labels_and_boxes):
+        """
+        Saves a single frames as well as it predicted annotation.
+        It should save 2 separate files under the same name, 1 .png for the raw frame and 1 .txt for the annotations.
+
+        Args:
+            frame: The current frame to be saved.
+            predicted_frames: Frames with predictions that might need to be saved alongside the original.
+            cv2: The OpenCV module used for image processing, passed in to avoid tight coupling.
+            labels_and_boxes: A list containing labels and their corresponding bounding boxes for the frame.
+        """
         pass

exports/iyolov8_export.py

@@ -2,15 +2,35 @@
 
 
 class IYolov8Export(ABC):
-
+    """
+    Interface for Yolov8 Export.
+    """
     @abstractmethod
     def initialize_save_dir(self):
+        """
+        Initializes save directory for Yolov8 export format
+        """
         pass
 
     @abstractmethod
     def save_frame(self, frame, predicted_frames, cv2, labels_and_boxes):
+        """
+        Saves a single frames as well as it predicted annotation.
+        It should save 2 separate files under the same name,
+            - 1 .png for the raw frame and is saved in images subdirectory.
+            - 1 .txt for the annotations and is saved in labels subdirectory.
+
+        Args:
+            frame: The current frame to be saved.
+            predicted_frames: Frames with predictions that might need to be saved alongside the original.
+            cv2: The OpenCV module used for image processing, passed in to avoid tight coupling.
+            labels_and_boxes: A list containing labels and their corresponding bounding boxes for the frame.
+        """
         pass
 
     @abstractmethod
     def create_yaml(self, model2):
+        """
+        Create .yaml file to map annotation labels with their corresponding names.
+        """
         pass

exports/yolov8_export.py

@@ -10,6 +10,11 @@
 
 
 class Yolov8Export(IYolov8Export):
+    """
+    Yolov8 Export class that implements functions for
+    initializing, saving frame and creating yaml file under specific format.
+    """
+
     def __init__(self, proj_dir_name):
         """
         Constructor.
@@ -25,10 +30,10 @@ def __init__(self, proj_dir_name):
 
     def initialize_save_dir(self):
         """
-        See ibase_project.py
+        See iyolov8_export.py
 
         Returns:
-            None
+            Success true or false.
         """
         self.result_dir_path = pjoin(self.proj_dir, f'{self._var.DATASET_FORMAT}-v{self._var.DATASET_VERSION}')
         os.makedirs(self.result_dir_path, exist_ok=True)
@@ -51,6 +56,12 @@ def initialize_save_dir(self):
             return False
 
     def save_frame(self, frame, predicted_frames, cv2, labels_and_boxes):
+        """
+        See iyolov8_export.py
+
+        Returns:
+            Predicted frame counter.
+        """
         print(f'5.1. Condition met, processing valid frame: {predicted_frames}')
         # Save original frame
         unix_time = int(time.time())

integrations/integration_factory.py

@@ -4,6 +4,9 @@
 
 
 class IntegrationFactory:
+    """
+    Integration Factory initializes specific integration types.
+    """
     def __init__(self):
         self._var = VariableClass()
         self.name = self._var.INTEGRATION_NAME

integrations/iroboflow_integration.py

@@ -2,11 +2,24 @@
 
 
 class IRoboflowIntegration(ABC):
+    """
+    Interface for Roboflow Integration class.
+    """
 
     @abstractmethod
     def upload_dataset(self, src_project_path):
+        """
+        Upload dataset to Roboflow platform.
+
+        Args:
+            src_project_path: Project save path
+        """
         pass
 
     @abstractmethod
     def __connect__(self):
+        """
+        Connect to Roboflow agent.
+        You need to provide Roboflow parameters in .env file.
+        """
         pass

integrations/is3_integration.py

@@ -2,19 +2,49 @@
 
 
 class IS3Integration(ABC):
+    """
+    Interface for S3 Integration class.
+    """
 
     @abstractmethod
     def upload_file(self, source_path, output_path):
+        """
+        Upload a single file to S3 compatible platform.
+
+        Args:
+            source_path: File save path
+            output_path: Desired path we want to save in S3
+        """
         pass
 
     @abstractmethod
     def upload_dataset(self, src_project_path):
+        """
+        Upload dataset to S3 compatible platform.
+
+        Args:
+            src_project_path: Projecet save path
+        """
         pass
 
     @abstractmethod
     def __connect__(self):
+        """
+        Connect to S3 compatible agent.
+        You need to provide S3 parameters in .env file.
+        """
         pass
 
     @abstractmethod
     def __check_bucket_exists__(self, bucket_name):
+        """
+        Check if input bucket exists after connecting to S3 compatible agent.
+        You need to provide S3 parameters in .env file.
+
+        Args:
+            bucket_name: Bucket name.
+
+        Returns:
+            True or False
+        """
         pass

integrations/roboflow_integration.py

@@ -7,11 +7,27 @@
 
 
 class RoboflowIntegration:
+    """
+    Roboflow Integration class that implements functions for connecting, uploading dataset
+    to Roboflow platform.
+    """
+
     def __init__(self):
+        """
+        Constructor.
+        """
         self._var = VariableClass()
         self.agent, self.ws, self.project = self.__connect__()
 
     def __connect__(self):
+        """
+        See iroboflow_integration.py
+
+        Returns:
+            agent: Connected agent.
+            workspace: Selected workspace in that agent.
+            project: Selected project in that workspace.
+        """
         try:
             # Attempt to initialize Roboflow with the API key
             agent = roboflow.Roboflow(api_key=self._var.ROBOFLOW_API_KEY)
@@ -29,15 +45,18 @@ def __connect__(self):
             raise ConnectionRefusedError(f'Error during Roboflow login: {e}')
 
     def upload_dataset(self, src_project_path):
+        """
+        See iroboflow_integration.py
+        """
         # Upload data set to an existing project
         self.ws.upload_dataset(
-                src_project_path,
-                pbasename(self.project.id),
-                num_workers=10,
-                project_license="MIT",
-                project_type="object-detection",
-                batch_name=None,
-                num_retries=0
+            src_project_path,
+            pbasename(self.project.id),
+            num_workers=10,
+            project_license="MIT",
+            project_type="object-detection",
+            batch_name=None,
+            num_retries=0
         )
         print('Uploaded')