File size: 2,952 Bytes
9fc21e6 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 |
# from ultralytics import YOLOv10
import torch
from config.set_config import set_configuration
set_config_project = set_configuration()
layout_model_weights_path = set_config_project.layout_model_weights_path
no_of_threads = set_config_project.no_of_threads
# def load_model_for_process(detection_model_path=layout_model_weights_path):
# """
# Load model in each subprocess to avoid CUDA initialization issues
# Returns:
# Model loaded in appropriate device
# """
# # Your model loading logic
# device = "cuda" if torch.cuda.is_available() else "cpu"
# # print(f"Using device: {device}")
# model = YOLOv10(detection_model_path).to(device)
# class_names = model.names
# class_names["11"] = "Table-header"
# class_names["12"] = "Portfolio-Company-Table"
# return model, class_names
import torch
from ultralytics import YOLO
# def load_model_for_process(detection_model_path=layout_model_weights_path):
# """
# Load model in each subprocess to avoid CUDA initialization issues
# Returns:
# Model loaded in appropriate device
# """
# # Your model loading logic
# device = "cuda" if torch.cuda.is_available() else "cpu"
# # print(f"Using device: {device}")
# model = YOLO(detection_model_path).to(device)
# class_names = model.names
# class_names["11"] = "Table-header"
# class_names["12"] = "Portfolio-Company-Table"
# print("YOLOV12"*10)
# return model, class_names
'''Below code for docling heron model'''
from transformers import RTDetrV2ForObjectDetection, RTDetrImageProcessor
# MODEL_NAME_DOCLING = "ds4sd/docling-layout-heron"
MODEL_NAME_DOCLING = layout_model_weights_path
def load_model_for_process(model_name=MODEL_NAME_DOCLING):
"""
Load the Docling Heron model and image processor in each subprocess to avoid CUDA initialization issues.
Returns:
Tuple of (model, image_processor, class_names)
"""
device = "cuda" if torch.cuda.is_available() else "cpu"
print(f"Using device: {device}")
# Load the image processor and model
image_processor = RTDetrImageProcessor.from_pretrained(model_name)
model = RTDetrV2ForObjectDetection.from_pretrained(model_name).to(device)
# Define class names mapping
class_names = {
0: "Caption",
1: "Footnote",
2: "Formula",
3: "List-item",
4: "Page-footer",
5: "Page-header",
6: "Picture",
7: "Section-header",
8: "Table",
9: "Text",
10: "Title",
11: "Document Index",
12: "Code",
13: "Checkbox-Selected",
14: "Checkbox-Unselected",
15: "Form",
16: "Key-Value Region",
# Additional classes for compatibility with existing pipeline
17 : "Table-header",
18 : "Portfolio-Company-Table"
}
return model, image_processor, class_names
|