Alessio Grancini commited on
Commit
d35aae8
·
verified ·
1 Parent(s): 89c89f9

Update image_segmenter.py

Browse files
Files changed (1) hide show
  1. image_segmenter.py +3 -1
image_segmenter.py CHANGED
@@ -40,7 +40,9 @@ class ImageSegmenter:
40
  # params
41
  objects_data = []
42
  image = image.copy()
43
- predictions = self.model.predict(image)
 
 
44
 
45
  cls_ids = predictions[0].boxes.cls.cpu().numpy()
46
  bounding_boxes = predictions[0].boxes.xyxy.int().cpu().numpy()
 
40
  # params
41
  objects_data = []
42
  image = image.copy()
43
+ #predictions = self.model.predict(image)
44
+ # Use the original image resolution (1512x1008) instead of default YOLO resize
45
+ predictions = self.model.predict(image, imgsz=(1512, 1008), conf=self.confidence_threshold)
46
 
47
  cls_ids = predictions[0].boxes.cls.cpu().numpy()
48
  bounding_boxes = predictions[0].boxes.xyxy.int().cpu().numpy()