Improve code based on pylint and black suggestions

detection/ml_detection.py

@@ -1,15 +1,21 @@
+"""
+Object Detection module
+"""
+
+import io
+import torch
 from transformers import DetrImageProcessor, DetrForObjectDetection
 from transformers import YolosImageProcessor, YolosForObjectDetection
-import torch
 from PIL import Image
-import io
 
 
 # Load transformer-based model (Yolos or DETR)
 def load_model(model_uri: str):
-    """Load Transformer model"""
-    """  - Doc DETR: https://huggingface.co/docs/transformers/en/model_doc/detr"""
-    """  - Doc Yolos: https://huggingface.co/docs/transformers/en/model_doc/yolos"""
+    """
+    Load Transformer model
+      - Doc DETR: https://huggingface.co/docs/transformers/en/model_doc/detr
+      - Doc Yolos: https://huggingface.co/docs/transformers/en/model_doc/yolos
+    """
 
     if "detr" in model_uri:
         # you can specify the revision tag if you don't want the timm dependency
@@ -27,9 +33,9 @@ def load_model(model_uri: str):
 def object_detection(processor, model, image_bytes):
     """Perform object detection task"""
 
-    print('Object detection prediction...')
-    #url = "http://images.cocodataset.org/val2017/000000039769.jpg"
-    #image = Image.open(requests.get(url, stream=True).raw)
+    print("Object detection prediction...")
+    # url = "http://images.cocodataset.org/val2017/000000039769.jpg"
+    # image = Image.open(requests.get(url, stream=True).raw)
 
     img = Image.open(io.BytesIO(image_bytes))
     inputs = processor(images=img, return_tensors="pt")
@@ -39,5 +45,7 @@ def object_detection(processor, model, image_bytes):
     # convert outputs (bounding boxes and class logits) to COCO API
     # let's only keep detections with score > 0.9
     target_sizes = torch.tensor([img.size[::-1]])
-    results = processor.post_process_object_detection(outputs, target_sizes=target_sizes, threshold=0.9)[0]
+    results = processor.post_process_object_detection(
+        outputs, target_sizes=target_sizes, threshold=0.9
+    )[0]
     return results

detection/ml_utils.py

@@ -1,17 +1,24 @@
+"""
+Utility functions for detection module
+"""
+
 import torch
-import json
 
 
 def convert_tensor_dict_to_json(tensor_dict):
     """Convert a dictionary of tensors to a JSON-serializable dictionary."""
 
-    json_dict = {key: value.detach().cpu().numpy().tolist() if isinstance(value, torch.Tensor) else value
-                  for key, value in tensor_dict.items()}
+    json_dict = {
+        key: (
+            value.detach().cpu().numpy().tolist()
+            if isinstance(value, torch.Tensor)
+            else value
+        )
+        for key, value in tensor_dict.items()
+    }
 
     # Convert to JSON string
     # json_str = json.dumps(json_dict)
     # print(json_str)
 
     return json_dict
-
-

inference_api.py

@@ -1,7 +1,11 @@
-import requests
-import argparse
+"""
+Object detection - command line inference via API
+"""
+
 import sys
 import base64
+import argparse
+import requests
 
 # Default examples
 # api = "http://localhost:8080/2015-03-31/functions/function/invocations"
@@ -12,15 +16,26 @@ def arg_parser():
     """Parse arguments"""
 
     # Create an ArgumentParser object
-    parser = argparse.ArgumentParser(description='Object detection inference via API call')
+    parser = argparse.ArgumentParser(
+        description="Object detection inference via API call"
+    )
     # Add arguments
-    parser.add_argument('--api', type=str, help='URL to server API (with endpoint)', required=True)
-    parser.add_argument('--file', type=str, help='Path to the input image file', required=True)
-    parser.add_argument('--model', type=str, \
-                        choices=['detr-resnet-50', 'detr-resnet-101', 'yolos-tiny', 'yolos-small'], \
-                        help='Model type', \
-                        required=False)
-    parser.add_argument('-v', '--verbose', action='store_true', help='Increase output verbosity')
+    parser.add_argument(
+        "--api", type=str, help="URL to server API (with endpoint)", required=True
+    )
+    parser.add_argument(
+        "--file", type=str, help="Path to the input image file", required=True
+    )
+    parser.add_argument(
+        "--model",
+        type=str,
+        choices=["detr-resnet-50", "detr-resnet-101", "yolos-tiny", "yolos-small"],
+        help="Model type",
+        required=False,
+    )
+    parser.add_argument(
+        "-v", "--verbose", action="store_true", help="Increase output verbosity"
+    )
     return parser
 
 
@@ -30,7 +45,7 @@ def main(args=None):
     args = arg_parser().parse_args(args)
     # Use the arguments
     if args.verbose:
-        print(f'Input file: {args.file}')
+        print(f"Input file: {args.file}")
 
     # Retrieve model type
     if args.model:
@@ -39,27 +54,27 @@ def main(args=None):
         model_name = ""
 
     # Load image
-    with open(args.file, 'rb') as image_file:
+    with open(args.file, "rb") as image_file:
         image_data = image_file.read()
 
     # Encode the image data in base64
-    encoded_image = base64.b64encode(image_data).decode('utf-8')
+    encoded_image = base64.b64encode(image_data).decode("utf-8")
 
     # Prepare the payload
     payload = {
-        'body': encoded_image,
-        'isBase64Encoded': True,
-        'model': model_name,
+        "body": encoded_image,
+        "isBase64Encoded": True,
+        "model": model_name,
     }
 
     # Send request to API
     # Option 'files': A dictionary of files to send to the specified url
     # response = requests.post(args.api, files={'image': image_data})
     # Option 'json': A JSON object to send to the specified url
-    response = requests.post(args.api, json = payload)
+    response = requests.post(args.api, json=payload, timeout=60)
 
     if response.status_code == 200:
-        print('Detection Results:')
+        print("Detection Results:")
         # Process the response
         # processed_data = json.loads(response.content)
         # print('processed_data', processed_data)

inference_direct.py

@@ -1,10 +1,13 @@
+"""
+Direct inference with hard-coded data
+"""
+
 from detection import ml_detection, ml_utils
-import json
-from PIL import Image
 
 
 # Run detection pipeline: load ML model, perform object detection and return json object
 def detection_pipeline(model_type, image_bytes):
+    """Detection pipeline: load ML model, perform object detection and return json object"""
     # Load correct ML model
     detr_processor, detr_model = ml_detection.load_model(model_type)
 
@@ -18,15 +21,16 @@ def detection_pipeline(model_type, image_bytes):
 
 
 def main():
-    print('Main function')
+    """Main function"""
+    print("Main function")
 
     model_type = "facebook/detr-resnet-50"
-    image_path = './samples/boats.jpg'
+    image_path = "./samples/boats.jpg"
 
     # Reading image file as image_bytes (similar to API request)
-    print('Reading image file...')
-    with open(image_path, 'rb') as image_file:
-       image_bytes = image_file.read()
+    print("Reading image file...")
+    with open(image_path, "rb") as image_file:
+        image_bytes = image_file.read()
 
     result_json = detection_pipeline(model_type, image_bytes)
     print("result_json:", result_json)
@@ -34,4 +38,3 @@ def main():
 
 if __name__ == "__main__":
     main()
-

lambda_function.py

@@ -1,7 +1,11 @@
-from detection import ml_detection, ml_utils
+"""
+AWS Lambda function
+"""
+
 import base64
 import json
 import logging
+from detection import ml_detection, ml_utils
 
 
 logger = logging.getLogger()
@@ -10,6 +14,7 @@ logger.setLevel(logging.INFO)
 
 # Find ML model type based on string request
 def get_model_type(query_string):
+    """Find ML model type based on string request"""
     # Default ml model type
     if query_string == "":
         model_type = "facebook/detr-resnet-50"
@@ -19,12 +24,13 @@ def get_model_type(query_string):
     elif "yolos" in query_string:
         model_type = "hustvl/" + query_string
     else:
-        raise Exception('Incorrect model type.')
+        raise Exception("Incorrect model type.")
     return model_type
 
 
 # Run detection pipeline: load ML model, perform object detection and return json object
 def detection_pipeline(model_type, image_bytes):
+    """detection pipeline: load ML model, perform object detection and return json object"""
     # Load correct ML model
     processor, model = ml_detection.load_model(model_type)
 
@@ -38,36 +44,43 @@ def detection_pipeline(model_type, image_bytes):
 
 
 def lambda_handler(event, context):
+    """
+        Lambda handler (proxy integration option unchecked on AWS API Gateway)
+
+        Args:
+            event (dict): The event that triggered the Lambda function.
+            context (LambdaContext): Information about the execution environment.
+
+        Returns:
+            dict: The response to be returned from the Lambda function.
+        """
+
     # logger.info(f"API event: {event}")
     try:
         # Retrieve model type
-        model_query = event.get('model', '')
+        model_query = event.get("model", "")
         model_type = get_model_type(model_query)
-        logger.info(f"Model query: {model_query}")
-        logger.info(f"Model type: {model_type}")
+        logger.info("Model query: %s", model_query)
+        logger.info("Model type: %s", model_type)
 
         # Decode the base64-encoded image data from the event
-        image_data = event['body']
-        if event['isBase64Encoded']:
+        image_data = event["body"]
+        if event["isBase64Encoded"]:
             image_data = base64.b64decode(image_data)
 
         # Run detection pipeline
         result_dict = detection_pipeline(model_type, image_data)
-        logger.info(f"API Results: {result_dict}")
+        logger.info("API Results: %s", str(result_dict))
 
         return {
-            'statusCode': 200,
-            'headers': {
-                'Content-Type': 'application/json'
-            },
-            'body': json.dumps(result_dict),
+            "statusCode": 200,
+            "headers": {"Content-Type": "application/json"},
+            "body": json.dumps(result_dict),
         }
     except Exception as e:
-        logger.info(f"API Error: {str(e)}")
+        logger.info("API Error: %s", str(e))
         return {
-            'statusCode': 500,
-            'headers': {
-                'Content-Type': 'application/json'
-            },
-            'body': json.dumps({'error': str(e)}),
-        }
+            "statusCode": 500,
+            "headers": {"Content-Type": "application/json"},
+            "body": json.dumps({"error": str(e)}),
+        }

tests/test_lambda.py

@@ -1,35 +1,41 @@
+"""
+Testing Lambda handler
+"""
+
 import os
 import sys
-import pytest
 import json
 import base64
+import pytest
+
+from lambda_function import lambda_handler
 
 
 current_dir = os.path.dirname(os.path.abspath(__file__))
 parent_dir = os.path.dirname(current_dir)
 sys.path.insert(0, os.path.dirname(parent_dir))
 
-from lambda_function import lambda_handler
-
 
 @pytest.fixture
 def event():
+    """Example json event"""
+
     # Get the directory of the current test file
     test_dir = os.path.dirname(os.path.abspath(__file__))
     # Construct the image path relative to the test directory
-    image_path = os.path.join(test_dir, 'data', 'savanna.jpg')
+    image_path = os.path.join(test_dir, "data", "savanna.jpg")
 
     # Read image data
-    with open(image_path, 'rb') as image_file:
+    with open(image_path, "rb") as image_file:
         image_data = image_file.read()
 
     # Encode the image data in base64
-    encoded_image = base64.b64encode(image_data).decode('utf-8')
+    encoded_image = base64.b64encode(image_data).decode("utf-8")
 
     # Prepare the payload
     json_event = {
-        'body': encoded_image,
-        'isBase64Encoded': True,
+        "body": encoded_image,
+        "isBase64Encoded": True,
     }
 
     return json_event
@@ -37,22 +43,25 @@ def event():
 
 @pytest.fixture
 def context():
+    """Example context"""
     return None
 
 
 def test_lambda_handler(event, context):
+    """Tests lambda handler"""
+
     lambda_response = lambda_handler(event, context)
     response_data = json.loads(lambda_response["body"])
 
-    print("lambda_response - type",type(lambda_response))
+    print("lambda_response - type", type(lambda_response))
     print("lambda_response", lambda_response)
     print("response_data - type", type(response_data))
     print("response_data", response_data)
 
     response_keys = list(response_data.keys())
-    gt_keys = ['scores', 'labels', 'boxes']
+    gt_keys = ["scores", "labels", "boxes"]
 
     assert lambda_response["statusCode"] == 200
     assert set(response_keys) == set(gt_keys), "Response keys do not match ground truth"
-    assert len(response_data['scores']) == 5
-    assert len(response_data['labels']) == 5
+    assert len(response_data["scores"]) == 5
+    assert len(response_data["labels"]) == 5

tests/test_ml_detection.py

@@ -1,31 +1,46 @@
-from detection import ml_detection
+"""
+Testing Detection module
+"""
+
 import os
 import pytest
 
+from detection import ml_detection
+
 
 # Test model loading
-@pytest.mark.parametrize("test_model_uri",[
-    ("facebook/detr-resnet-50"),
-    ("facebook/detr-resnet-101"),
-])
+@pytest.mark.parametrize(
+    "test_model_uri",
+    [
+        ("facebook/detr-resnet-50"),
+        ("facebook/detr-resnet-101"),
+    ],
+)
 def test_load_model(test_model_uri):
+    """Testing model loading"""
+
     processor, model = ml_detection.load_model(test_model_uri)
     assert processor is not None
     assert model is not None
 
 
 # Test image detection
-@pytest.mark.parametrize("test_model_uri", [
-    ("facebook/detr-resnet-50"),
-    ("facebook/detr-resnet-101"),
-])
+@pytest.mark.parametrize(
+    "test_model_uri",
+    [
+        ("facebook/detr-resnet-50"),
+        ("facebook/detr-resnet-101"),
+    ],
+)
 def test_object_detection(test_model_uri):
+    """Testing object detection function"""
+
     processor, model = ml_detection.load_model(test_model_uri)
 
     # Get the directory of the current test file
     test_dir = os.path.dirname(os.path.abspath(__file__))
     # Construct the image path relative to the test directory
-    image_path = os.path.join(test_dir, 'data', 'savanna.jpg')
+    image_path = os.path.join(test_dir, "data", "savanna.jpg")
 
     with open(image_path, "rb") as f:
         image_bytes = f.read()

tests/test_ml_utils.py

@@ -1,11 +1,16 @@
-from detection import ml_utils
+"""
+Testing Detection utility functions
+"""
+
 import torch
-import json
+from detection import ml_utils
 
 
 # Test dictionary conversion
 def test_convert_tensor_dict_to_json():
-    my_dict = {'scores': torch.tensor([1, 2, 3])}
-    my_list_gt = {'scores': [1, 2, 3]}
+    """Testing tensor to dict conversions"""
+
+    my_dict = {"scores": torch.tensor([1, 2, 3])}
+    my_list_gt = {"scores": [1, 2, 3]}
     my_list = ml_utils.convert_tensor_dict_to_json(my_dict)
     assert my_list == my_list_gt