m7mdal7aj commited on
Commit
4084ed7
1 Parent(s): 7537a6a

Update my_model/object_detection.py

Browse files
Files changed (1) hide show
  1. my_model/object_detection.py +12 -9
my_model/object_detection.py CHANGED
@@ -87,27 +87,30 @@ class ObjectDetector:
87
  print(f"Error loading YOLOv5 model: {e}")
88
  raise
89
 
90
- def process_image(self, image_path):
91
  """
92
- Process the image from the given path.
93
-
94
  Args:
95
- image_path (str): Path to the image file.
96
-
97
  Returns:
98
  Image.Image: Processed image in RGB format.
99
-
100
  Raises:
101
  Exception: If an error occurs during image processing.
102
  """
103
-
104
  try:
105
- with Image.open(image_path) as image:
106
- return image.convert("RGB")
 
 
 
 
 
 
107
  except Exception as e:
108
  print(f"Error processing image: {e}")
109
  raise
110
 
 
111
  def detect_objects(self, image, threshold=0.4):
112
  """
113
  Detect objects in the given image using the loaded model.
 
87
  print(f"Error loading YOLOv5 model: {e}")
88
  raise
89
 
90
+ def process_image(self, image_input):
91
  """
92
+ Process the image from the given path or file-like object.
 
93
  Args:
94
+ image_input (str or file-like object): Path to the image file or a file-like object.
 
95
  Returns:
96
  Image.Image: Processed image in RGB format.
 
97
  Raises:
98
  Exception: If an error occurs during image processing.
99
  """
 
100
  try:
101
+ # Check if the input is a string (path) or a file-like object
102
+ if isinstance(image_input, str):
103
+ # Open the image from a file path
104
+ with Image.open(image_input) as image:
105
+ return image.convert("RGB")
106
+ else:
107
+ # Assume image_input is a file-like object and process it directly
108
+ return Image.open(image_input).convert("RGB")
109
  except Exception as e:
110
  print(f"Error processing image: {e}")
111
  raise
112
 
113
+
114
  def detect_objects(self, image, threshold=0.4):
115
  """
116
  Detect objects in the given image using the loaded model.