Cleaning the code

Author: javanese-programmer

capture.py

@@ -1,98 +1,122 @@
-"""Script to capture images from object detection"""
+"""Script to capture images from object detection."""
 
 import argparse
-from utils import video
+from utils import video, terminal
 from detection import deeplearning, traditional
 
 
 def parse_argument():
-    """Parse arguments for command line"""
+    """Parse arguments for command line."""
     parser = argparse.ArgumentParser(
-        formatter_class=argparse.ArgumentDefaultsHelpFormatter)
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter
+    )
     parser.add_argument(
-        '--method',
-        help='Two computer vision method: traditional or deeplearning',
+        "--method",
+        help="Two computer vision method: traditional or deeplearning",
         required=False,
         type=str,
-        default='deeplearning')
+        default="deeplearning",
+    )
     parser.add_argument(
-        '--detectionType',
-        help='Type of detection for traditional method: color OR shape',
+        "--detectionType",
+        help="Type of detection for traditional method: color OR shape",
         required=False,
-        default='color')
+        default="color",
+    )
     parser.add_argument(
-        '--multipleObject',
-        help='Whether to detect single or multiple types of objects.',
+        "--multipleObject",
+        help="Whether to detect single or multiple types of objects.",
         required=False,
-        default=False)
+        default=False,
+    )
     parser.add_argument(
-        '--trueClass',
-        help='Name of the object class to be detected if it is a single type.',
+        "--trueClass",
+        help="Name of the object class to be detected if it is a single type.",
         required=False,
         type=str,
-        default='yellow_duck')
+        default="yellow_duck",
+    )
     parser.add_argument(
-        '--collectAll',
-        help='Whether to collect all images or only the True Positive one.',
+        "--collectAll",
+        help="Whether to collect all images or only the True Positive one.",
         required=False,
-        default=False)
+        default=False,
+    )
     parser.add_argument(
-        '--model',
-        help='Path of the object detection model.',
+        "--model",
+        help="Path of the object detection model.",
         required=False,
-        default='./model/frogducky2.tflite')
+        default="./model/frogducky2.tflite",
+    )
     parser.add_argument(
-        '--frameWidth',
-        help='Width of frame to capture from camera.',
+        "--frameWidth",
+        help="Width of frame to capture from camera.",
         required=False,
         type=int,
-        default=640)
+        default=640,
+    )
     parser.add_argument(
-        '--frameHeight',
-        help='Height of frame to capture from camera.',
+        "--frameHeight",
+        help="Height of frame to capture from camera.",
         required=False,
         type=int,
-        default=480)
+        default=480,
+    )
     parser.add_argument(
-        '--numThreads',
-        help='Number of CPU threads to run the model.',
+        "--numThreads",
+        help="Number of CPU threads to run the model.",
         required=False,
         type=int,
-        default=4)
+        default=4,
+    )
     parser.add_argument(
-        '--enableEdgeTPU',
-        help='Whether to run the model on EdgeTPU.',
-        action='store_true',
+        "--enableEdgeTPU",
+        help="Whether to run the model on EdgeTPU.",
+        action="store_true",
         required=False,
-        default=False)
+        default=False,
+    )
     parser.add_argument(
-        '--videoFilename',
-        help='Name of the output video file with .mp4 extension',
+        "--videoFilename",
+        help="Name of the output video file with .mp4 extension",
         required=False,
-        default='Deteksi_Objek.mp4')
+        default="Deteksi_Objek.mp4",
+    )
     return parser.parse_args()
 
 
 def main():
-    """Main function to run detection"""
+    """Run detection and capture the result."""
     args = parse_argument()
-    
-    if args.method == 'deeplearning':
-        detector = deeplearning.DeepDetector(bool(args.multipleObject), args.model,
-                                                 args.frameWidth, args.frameHeight,
-                                                 int(args.numThreads),bool(args.enableEdgeTPU))
-        detector.capture(str(args.trueClass), bool(args.collectAll), str(args.videoFilename))
-    
-    elif args.method == 'traditional':   
-        if str(args.detectionType) == 'category':
+
+    if args.method == "deeplearning":
+        detector = deeplearning.DeepDetector(
+            bool(args.multipleObject),
+            args.model,
+            args.frameWidth,
+            args.frameHeight,
+            int(args.numThreads),
+            bool(args.enableEdgeTPU),
+        )
+        detector.capture(
+            str(args.trueClass), bool(args.collectAll), str(args.videoFilename)
+        )
+
+    elif args.method == "traditional":
+        if str(args.detectionType) == "category":
             args.detectionType = terminal.prompt_type()
-        detector = traditional.TraditionalDetector(bool(args.multipleObject), args.frameWidth,
-                                                   args.frameHeight)
-        detector.capture(str(args.detectionType), str(args.trueClass), 
-                         bool(args.collectAll), str(args.videoFilename))
-        
+        detector = traditional.TraditionalDetector(
+            bool(args.multipleObject), args.frameWidth, args.frameHeight
+        )
+        detector.capture(
+            str(args.detectionType),
+            str(args.trueClass),
+            bool(args.collectAll),
+            str(args.videoFilename),
+        )
+
     video.move_video(str(args.videoFilename))
 
 
-if __name__ == '__main__':
-    main()
+if __name__ == "__main__":
+    main()

collect.py

@@ -1,7 +1,7 @@
-import os
+"""Scripts to collect images without running a detection."""
+
 import sys
 
-import numpy as np
 import cv2
 
 IMG_ADDR = "/home/pi/Desktop/Object Detection/Raspberry Pi"
@@ -17,19 +17,20 @@
     success, image = cap.read()
     if not success:
         sys.exit(
-            'ERROR: Unable to read from webcam. Please verify your webcam settings.'
+            """ERROR: Unable to read from webcam.
+            Please verify your webcam settings."""
         )
-    
+
     counter += 1
-    
-    #image = cv2.flip(image, 1)
-#     image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+
+    # image = cv2.flip(image, 1)
+    #     image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
     cv2.imwrite(IMG_ADDR + "/image/Image" + str(counter) + ".png", image)
 
     # Stop the program if the ESC key is pressed.
     if cv2.waitKey(1) == 27:
         break
-    
+
     cv2.imshow("Captured Image", image)
-    
-cv2.destroyWindow("Captured Image")
+
+cv2.destroyWindow("Captured Image")