Update my_model/object_detection.py

my_model/object_detection.py

@@ -8,8 +8,6 @@ import os
 from my_model.gen_utilities import get_image_path, get_model_path ,show_image
 
 
-
-
 class ObjectDetector:
     """
         A class for detecting objects in images using models like Detic and YOLOv5.
@@ -63,7 +61,6 @@ class ObjectDetector:
 
         try:
             model_path = get_model_path('deformable-detr-detic')
-           
             self.processor = AutoImageProcessor.from_pretrained(model_path)
             self.model = AutoModelForObjectDetection.from_pretrained(model_path)
         except Exception as e:
@@ -115,8 +112,7 @@ class ObjectDetector:
             print(f"Error processing image: {e}")
             raise
 
-
-
+    
     def detect_objects(self, image, threshold=0.4):
         """
         Detect objects in the given image using the loaded model.
@@ -139,6 +135,7 @@ class ObjectDetector:
         else:
             raise ValueError("Model not loaded or unsupported model name")
 
+    
     def _detect_with_detic(self, image, threshold):
         """
         Detect objects using the Detic model.
@@ -155,9 +152,7 @@ class ObjectDetector:
         inputs = self.processor(images=image, return_tensors="pt")
         outputs = self.model(**inputs)
         target_sizes = torch.tensor([image.size[::-1]])
-        results = self.processor.post_process_object_detection(outputs, target_sizes=target_sizes, threshold=threshold)[
-            0]
-
+        results = self.processor.post_process_object_detection(outputs, target_sizes=target_sizes, threshold=threshold)[0]
         detected_objects_str = ""
         detected_objects_list = []
         for score, label, box in zip(results["scores"], results["labels"], results["boxes"]):
@@ -169,6 +164,7 @@ class ObjectDetector:
                 detected_objects_list.append((label_name, box_rounded, certainty))
         return detected_objects_str, detected_objects_list
 
+    
     def _detect_with_yolov5(self, image, threshold):
         """
         Detect objects using the YOLOv5 model.
@@ -184,7 +180,6 @@ class ObjectDetector:
 
         cv2_img = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
         results = self.model(cv2_img)
-
         detected_objects_str = ""
         detected_objects_list = []
         for *bbox, conf, cls in results.xyxy[0]:
@@ -214,7 +209,6 @@ class ObjectDetector:
             font = ImageFont.truetype("arial.ttf", 15)
         except IOError:
             font = ImageFont.load_default()
-
         colors = ["red", "green", "blue", "yellow", "purple", "orange"]
         label_color_map = {}
 
@@ -224,7 +218,6 @@ class ObjectDetector:
 
             color = label_color_map[label_name]
             draw.rectangle(box, outline=color, width=3)
-
             label_text = f"{label_name}"
             if show_confidence:
                 label_text += f" ({round(score, 2)}%)"