|
|
@@ -314,7 +314,7 @@ class CustomPEKModel:
|
|
|
elif self.layout_model_name == MODEL_NAME.DocLayout_YOLO:
|
|
|
# doclayout_yolo
|
|
|
layout_res = []
|
|
|
- doclayout_yolo_res = self.layout_model.predict(image, imgsz=1024, conf=0.15, iou=0.45, verbose=True, device=self.device)[0]
|
|
|
+ doclayout_yolo_res = self.layout_model.predict(image, imgsz=1024, conf=0.25, iou=0.45, verbose=True, device=self.device)[0]
|
|
|
for xyxy, conf, cla in zip(doclayout_yolo_res.boxes.xyxy.cpu(), doclayout_yolo_res.boxes.conf.cpu(), doclayout_yolo_res.boxes.cls.cpu()):
|
|
|
xmin, ymin, xmax, ymax = [int(p.item()) for p in xyxy]
|
|
|
new_item = {
|