Update my_model/object_detection.py

my_model/object_detection.py

@@ -26,10 +26,11 @@ class ObjectDetector:
         """
         Initializes the ObjectDetector class with default values.
         """
-
+        
         self.model = None
         self.processor = None
         self.model_name = None
+        
 
     def load_model(self, model_name='detic', pretrained=True, model_version='yolov5s'):
         """
@@ -52,6 +53,7 @@ class ObjectDetector:
         else:
             raise ValueError(f"Unsupported model name: {model_name}")
 
+    
     def _load_detic_model(self, pretrained):
         """
         Load the Detic model.
@@ -62,13 +64,13 @@ class ObjectDetector:
 
         try:
             model_path = get_model_path('deformable-detr-detic')
-            st.write(model_path)
             self.processor = AutoImageProcessor.from_pretrained(model_path)
             self.model = AutoModelForObjectDetection.from_pretrained(model_path)
         except Exception as e:
             print(f"Error loading Detic model: {e}")
             raise
 
+    
     def _load_yolov5_model(self, pretrained, model_version):
         """
         Load the YOLOv5 model.
@@ -80,7 +82,6 @@ class ObjectDetector:
 
         try:
             model_path = get_model_path ('yolov5')
-            st.write(model_path)
             if model_path and os.path.exists(model_path):
                 self.model = torch.hub.load(model_path, model_version, pretrained=pretrained, source='local')
             else:
@@ -89,6 +90,7 @@ class ObjectDetector:
             print(f"Error loading YOLOv5 model: {e}")
             raise
 
+    
     def process_image(self, image_input):
         """
         Process the image from the given path or file-like object.
@@ -194,6 +196,7 @@ class ObjectDetector:
                 detected_objects_list.append((label_name, box_rounded, certainty))
         return detected_objects_str, detected_objects_list
 
+    
     def draw_boxes(self, image, detected_objects, show_confidence=True):
         """
         Draw bounding boxes around detected objects in the image.
@@ -218,7 +221,6 @@ class ObjectDetector:
         for label_name, box, score in detected_objects:
             if label_name not in label_color_map:
                 label_color_map[label_name] = colors[len(label_color_map) % len(colors)]
-
             color = label_color_map[label_name]
             draw.rectangle(box, outline=color, width=3)
             label_text = f"{label_name}"