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

Update utils.py

Browse files
Files changed (1) hide show
  1. utils.py +15 -7
utils.py CHANGED
@@ -374,13 +374,21 @@ def predict(model, image, caption, box_threshold, text_threshold):
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):
 
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):