Project cleanup (#174)

* Add linting and formatting checks

* remove unused imports

Author: Sarah5567

scripts/export_to_onnx.py

@@ -1,18 +1,17 @@
-import os
-import sys
 import argparse
 from pathlib import Path
 
 import torch
 
-from src.inference.utils.inference_factory import InferenceFactory, InferenceConfig, Backend, ModelArch
+from src.inference.utils.inference_factory import Backend, InferenceConfig, InferenceFactory, ModelArch
 from src.path_utils import ensure_clean_directory
 
 NUM_CLASSES = 83
 INPUT_SIZE = (256, 256)
 DEVICE = "cpu"
 MODELS_DIR_PATH = Path("models")
 
+
 def main(model_name: str):
     arch = ModelArch.RESNET if "resnet" in model_name.lower() else ModelArch.MOBILENET
     backend = Backend.PYTORCH

scripts/onnx_validation.py

@@ -1,15 +1,14 @@
 import os
-import argparse
 from pathlib import Path
 
 import numpy as np
 import torch
 from PIL import Image
 
 from dataset.optimal_class_mapping import MODEL_NAMES as ID_TO_NAME
-from src.inference.utils.inference_factory import InferenceFactory, InferenceConfig, Backend, ModelArch
 from src.inference.base.classifier_inference_base import ClassifierInferenceBase
 from src.inference.base.classifier_inference_base_onnx import OnnxClassifierInferenceBase
+from src.inference.utils.inference_factory import Backend, InferenceConfig, InferenceFactory, ModelArch
 
 NUM_CLASSES = 83
 INPUT_SIZE = (256, 256)

services/consumer/consumer.py

@@ -1,28 +1,32 @@
-import paho.mqtt.client as mqtt
-import pymongo
 import json
 from datetime import datetime
 
+import paho.mqtt.client as mqtt
+import pymongo
+
+
 def on_connect(client, userdata, flags, rc):
     print(f"🔗 Connected: {rc}")
     client.subscribe("deepstream/predictions")
 
+
 def on_message(client, userdata, msg):
     try:
         payload = json.loads(msg.payload.decode())
         payload["received_at"] = datetime.now().isoformat()
-        
+
         mongo_client = pymongo.MongoClient("mongodb://agstream_mongo:27017/")
         db = mongo_client["agstream"]
         collection = db["predictions"]
-        
+
         collection.insert_one(payload)
-        print(f"💾 Saved to MongoDB!")
+        print("💾 Saved to MongoDB!")
         mongo_client.close()
-        
+
     except Exception as e:
         print(f"❌ Error: {e}")
 
+
 client = mqtt.Client()
 client.on_connect = on_connect
 client.on_message = on_message

services/consumer/mqtt_consumer.py

@@ -1,59 +1,60 @@
 #!/usr/bin/env python3
-import paho.mqtt.client as mqtt
-import pymongo
 import json
 from datetime import datetime
 
+import paho.mqtt.client as mqtt
+import pymongo
+
 MQTT_BROKER = "agstream_mosquitto"
 MQTT_PORT = 1883
 MQTT_TOPIC = "deepstream/predictions"
 MONGO_URI = "mongodb://agstream_mongo:27017/"
 MONGO_DB = "agstream"
 MONGO_COLLECTION = "predictions"
 
+
 def on_connect(client, userdata, flags, rc):
     print(f"✅ Connected to MQTT broker with result code {rc}")
     client.subscribe(MQTT_TOPIC)
 
+
 def on_message(client, userdata, msg):
     try:
         payload = json.loads(msg.payload.decode())
         payload["received_at"] = datetime.now().isoformat()
-        
+
         # Extract classification data from nvmsgbroker format
         obj = payload.get("object", {})
         if obj.get("id") != "0" and obj.get("id"):
             class_id = int(obj["id"])
             confidence = obj.get("confidence", 0)
-            
+
             # Add extracted classification
-            payload["classification"] = {
-                "class_id": class_id,
-                "confidence": confidence
-            }
+            payload["classification"] = {"class_id": class_id, "confidence": confidence}
             print(f"🌱 FOUND CLASSIFICATION: ID {class_id}, confidence {confidence:.3f}")
-        
+
         mongo_client = pymongo.MongoClient(MONGO_URI)
         db = mongo_client[MONGO_DB]
         collection = db[MONGO_COLLECTION]
-        
+
         result = collection.insert_one(payload)
-        
+
         if "classification" in payload:
-            print(f"✅ Saved classification to MongoDB!")
+            print("✅ Saved classification to MongoDB!")
         else:
-            print(f"✅ Saved: No classification")
-        
+            print("✅ Saved: No classification")
+
         mongo_client.close()
-        
+
     except Exception as e:
         print(f"❌ Error: {e}")
 
+
 if __name__ == "__main__":
     client = mqtt.Client()
     client.on_connect = on_connect
     client.on_message = on_message
-    
+
     print("🚀 Starting Enhanced MQTT Consumer...")
     client.connect(MQTT_BROKER, MQTT_PORT, 60)
     client.loop_forever()

src/deepstream/helpers/load_class_labels.py

@@ -2,6 +2,7 @@
 
 PLANT_LABELS = "/workspace/configs/crop_and_weed_83_classes.txt"
 
+
 # Load class labels
 def load_class_labels() -> List[str]:
     try:

src/deepstream/helpers/meta_tensor_extractor.py

@@ -1,8 +1,10 @@
 import ctypes
-import numpy.typing as npt
+
 import numpy as np
+import numpy.typing as npt
 import pyds
 
+
 class TensorExtractor:
     def extract_logits(self, tensor_meta) -> npt.NDArray[np.float32]:
         """
@@ -13,10 +15,7 @@ def extract_logits(self, tensor_meta) -> npt.NDArray[np.float32]:
         dims = [layer.dims.d[i] for i in range(layer.dims.numDims)]
         numel = int(np.prod(dims))
 
-        ptr = ctypes.cast(
-            pyds.get_ptr(layer.buffer),
-            ctypes.POINTER(ctypes.c_float)
-        )
+        ptr = ctypes.cast(pyds.get_ptr(layer.buffer), ctypes.POINTER(ctypes.c_float))
         logits = np.ctypeslib.as_array(ptr, shape=(numel,))
 
         # Copy so we are not tied to DeepStream's memory lifetime

src/deepstream/helpers/plant_msg_meta_builder.py

@@ -1,14 +1,12 @@
 import sys
-from pydantic import BaseModel
 
 import gi
-gi.require_version("Gst", "1.0")
-
-from gi.repository import Gst
 import pyds
+from pydantic import BaseModel
 
 from src.deepstream.helpers.softmax_topk_classifier import ClassificationPrediction
 
+
 class PlantEvent(BaseModel):
     frame_id: int
     plant_id: str

src/deepstream/helpers/remove_background.py

@@ -1,8 +1,9 @@
-import numpy as np
 import cv2
+import numpy as np
 
 MORPH_KERNEL = 34
 
+
 def remove_background(frame_bgr: np.ndarray) -> np.ndarray:
     rgb = cv2.cvtColor(frame_bgr, cv2.COLOR_BGR2RGB).astype(np.float32)
     R, G, B = rgb[..., 0], rgb[..., 1], rgb[..., 2]
@@ -12,10 +13,7 @@ def remove_background(frame_bgr: np.ndarray) -> np.ndarray:
 
     _, mask = cv2.threshold(exg_norm, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
 
-    kernel = cv2.getStructuringElement(
-        cv2.MORPH_ELLIPSE,
-        (MORPH_KERNEL, MORPH_KERNEL)
-    )
+    kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (MORPH_KERNEL, MORPH_KERNEL))
     mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, kernel)
 
     out = np.zeros_like(frame_bgr)

src/deepstream/helpers/should_skip_frame.py

@@ -1,19 +1,22 @@
-from typing import Any
 from enum import Enum
+from typing import Any
+
+import cv2
 import gi
 import numpy as np
 import pyds
-import cv2
 
 gi.require_version("Gst", "1.0")
 from gi.repository import Gst
 
 from src.frame_comparison.frame_change_detector import FrameChangeDetector
 
+
 class FrameProcessDecision(str, Enum):
     PROCESS = "process"
     SKIP = "skip"
 
+
 def should_skip_frame(info: Any, frame_meta: Any, batch_meta: Any, frame_change_detector: FrameChangeDetector) -> int:
     """Pad probe to drop frames based on frame difference analysis."""
     gst_buffer = info.get_buffer()
@@ -25,12 +28,8 @@ def should_skip_frame(info: Any, frame_meta: Any, batch_meta: Any, frame_change_
     frame_bgr = cv2.cvtColor(frame, cv2.COLOR_RGBA2BGR)
 
     should_process, metrics = frame_change_detector.should_process(frame_bgr)
-    
-    decision = (
-        FrameProcessDecision.PROCESS
-        if should_process
-        else FrameProcessDecision.SKIP
-    )
+
+    decision = FrameProcessDecision.PROCESS if should_process else FrameProcessDecision.SKIP
 
     print(
         f"Frame {frame_meta.frame_num:05d}: {decision.value} "

src/deepstream/helpers/softmax_topk_classifier.py

@@ -1,5 +1,6 @@
-import numpy as np
 from typing import List
+
+import numpy as np
 from pydantic import BaseModel
 
 
@@ -19,7 +20,7 @@ def predict_from_logits(self, logits: np.ndarray) -> List[ClassificationPredicti
         exp = np.exp(logits - np.max(logits))
         probs = exp / np.sum(exp)
 
-        top_idx = np.argsort(probs)[-self.top_k:][::-1]
+        top_idx = np.argsort(probs)[-self.top_k :][::-1]
 
         results: List[ClassificationPrediction] = []
         for idx in top_idx: