Update detection_ocr.py

detection_ocr.py

@@ -1,7 +1,6 @@
 import torch
 from PIL import Image, ImageDraw
 from typing import List, Dict, Any
-from huggingface_hub import hf_hub_download
 from doclayout_yolo import YOLOv10
 from pix2tex import cli as pix2tex
 from pathlib import Path
@@ -14,20 +13,22 @@ class DetectionOCR:
 
     def __init__(self, device=None):
         if device is None:
-            device = "cuda:0" if torch.cuda.is_available() else "cpu"
+            device = "cuda" if torch.cuda.is_available() else "cpu"
         self.device = device
         print(f"[DetectionOCR] Using {self.device}")
 
-        # Load YOLO model
-        ckpt = "yolo.pt"
-        print("[DetectionOCR] Loading YOLO checkpoint...")
-        self.yolo = YOLOv10(ckpt).to(self.device)
-        self.names = self.yolo.names  # ← FIX
+        # ---- YOLO (HF from_pretrained) ----
+        print("[DetectionOCR] Loading DocLayout-YOLO...")
+        self.yolo = YOLOv10.from_pretrained(
+            "juliozhao/DocLayout-YOLO-DocStructBench"
+        ).to(self.device)
+        self.names = self.yolo.names
 
-        # Load Pix2TeX OCR
-        print("[DetectionOCR] Loading LatexOCR...")
+        # ---- Pix2TeX ----
+        print("[DetectionOCR] Loading pix2tex OCR...")
         self.pix2tex = pix2tex.LatexOCR()
 
+
     # ------------------------------------------------------------
     # 1. DETECTION
     # ------------------------------------------------------------