banao-tech commited on
Commit
8c51f84
·
verified ·
1 Parent(s): 0fb0aae

Update utils.py

Browse files
Files changed (1) hide show
  1. utils.py +7 -15
utils.py CHANGED
@@ -374,21 +374,13 @@ def predict(model, image, caption, box_threshold, text_threshold):
374
  return boxes, logits, phrases
375
 
376
 
377
- def predict_yolo(model, image_path, box_threshold, imgsz, scale_img, iou_threshold=0.7):
378
- """Use YOLO model for object detection with correct parameters"""
379
- kwargs = {
380
- 'conf': box_threshold, # Correct confidence parameter
381
- 'iou': iou_threshold, # Correct IoU parameter
382
- 'verbose': False
383
- }
384
-
385
- if scale_img:
386
- kwargs['imgsz'] = imgsz
387
-
388
- results = model.predict(image_path, **kwargs)
389
- boxes = results[0].boxes.xyxy
390
- conf = results[0].boxes.conf
391
- return boxes, conf, [str(i) for i in range(len(boxes))]
392
 
393
 
394
  def get_som_labeled_img(img_path, model=None, BOX_TRESHOLD = 0.01, output_coord_in_ratio=False, ocr_bbox=None, text_scale=0.4, text_padding=5, draw_bbox_config=None, caption_model_processor=None, ocr_text=[], use_local_semantics=True, iou_threshold=0.9,prompt=None, scale_img=False, imgsz=None, batch_size=None):
 
374
  return boxes, logits, phrases
375
 
376
 
377
+ # utils.py (correct implementation)
378
+ def predict_yolo(model, image_path, box_threshold):
379
+ result = model.predict(source=image_path, conf=box_threshold) # ✅ Valid for YOLO
380
+ boxes = result[0].boxes.xyxy
381
+ conf = result[0].boxes.conf
382
+ phrases = [str(i) for i in range(len(boxes))]
383
+ return boxes, conf, phrases
 
 
 
 
 
 
 
 
384
 
385
 
386
  def get_som_labeled_img(img_path, model=None, BOX_TRESHOLD = 0.01, output_coord_in_ratio=False, ocr_bbox=None, text_scale=0.4, text_padding=5, draw_bbox_config=None, caption_model_processor=None, ocr_text=[], use_local_semantics=True, iou_threshold=0.9,prompt=None, scale_img=False, imgsz=None, batch_size=None):