Commit 27281c92 authored by myhloli's avatar myhloli
Browse files

refactor(magic_pdf): remove unnecessary half() calls for CPU devices

- Remove half() calls for DocLayoutYOLO and YOLOv8 models
- This change prevents potential errors when running models on CPU
parent 25286669
...@@ -4,8 +4,6 @@ from doclayout_yolo import YOLOv10 ...@@ -4,8 +4,6 @@ from doclayout_yolo import YOLOv10
class DocLayoutYOLOModel(object): class DocLayoutYOLOModel(object):
def __init__(self, weight, device): def __init__(self, weight, device):
self.model = YOLOv10(weight) self.model = YOLOv10(weight)
if not device.startswith("cpu"):
self.model.half()
self.device = device self.device = device
def predict(self, image): def predict(self, image):
......
...@@ -4,8 +4,6 @@ from ultralytics import YOLO ...@@ -4,8 +4,6 @@ from ultralytics import YOLO
class YOLOv8MFDModel(object): class YOLOv8MFDModel(object):
def __init__(self, weight, device="cpu"): def __init__(self, weight, device="cpu"):
self.mfd_model = YOLO(weight) self.mfd_model = YOLO(weight)
if not device.startswith("cpu"):
self.mfd_model.half()
self.device = device self.device = device
def predict(self, image): def predict(self, image):
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment