|
|
import modal |
|
|
from io import BytesIO |
|
|
from pathlib import Path |
|
|
from fastapi import File, UploadFile, Form |
|
|
image = ( |
|
|
modal.Image.debian_slim(python_version="3.10") |
|
|
.apt_install(["libgl1-mesa-glx", "libglib2.0-0"]) |
|
|
.pip_install( |
|
|
"ultralytics>=8.2.85", |
|
|
"doclayout-yolo==0.0.2", |
|
|
"huggingface-hub", |
|
|
"fastapi", |
|
|
) |
|
|
) |
|
|
volume = modal.Volume.from_name("yolo-layout-detection", create_if_missing=True) |
|
|
volume_path = Path("/root") / "data" |
|
|
model_path = volume_path / "path2doclayout_yolo_ft.pt" |
|
|
app = modal.App( |
|
|
"yolo-layout-detection-temp", |
|
|
image=image, |
|
|
volumes={volume_path: volume}, |
|
|
) |
|
|
@app.function() |
|
|
def download_model(): |
|
|
from huggingface_hub import snapshot_download |
|
|
snapshot_download( |
|
|
repo_id="opendatalab/pdf-extract-kit-1.0", |
|
|
local_dir=volume_path, |
|
|
allow_patterns='path2*', |
|
|
max_workers=20, |
|
|
) |
|
|
@app.cls(gpu="a10g") |
|
|
class LayoutDetection: |
|
|
@modal.enter() |
|
|
def load_model(self): |
|
|
from doclayout_yolo import YOLOv10 |
|
|
self.model = YOLOv10(model_path) |
|
|
@modal.web_endpoint(method="POST", docs=True) |
|
|
async def predict(self, img: UploadFile = File(...), task: str = Form(...)): |
|
|
from PIL import Image |
|
|
img_bytes = await img.read() |
|
|
img = Image.open(BytesIO(img_bytes)) |
|
|
results = self.model.predict(img) |
|
|
|
|
|
figs = [] |
|
|
for result in results: |
|
|
boxes = result.__dict__['boxes'].xyxy.cpu().tolist() |
|
|
classes = result.__dict__['boxes'].cls.cpu().tolist() |
|
|
scores = result.__dict__['boxes'].conf.cpu().tolist() |
|
|
targets, captions = [], [] |
|
|
for box, cls, score in zip(boxes, classes, scores): |
|
|
if task == "figure": |
|
|
if cls == 3: |
|
|
targets.append({"box": box, "score": score}) |
|
|
elif task == "table": |
|
|
if cls == 5: |
|
|
targets.append({"box": box, "score": score}) |
|
|
elif task == "figurecaption": |
|
|
if cls == 3: |
|
|
targets.append({"box": box, "score": score}) |
|
|
elif cls == 4: |
|
|
captions.append({"box": box, "score": score}) |
|
|
elif task == "tablecaption": |
|
|
if cls == 5: |
|
|
targets.append({"box": box, "score": score}) |
|
|
elif cls == 6 or cls == 7: |
|
|
captions.append({"box": box, "score": score}) |
|
|
if not captions: |
|
|
figs = targets |
|
|
else: |
|
|
matches = [] |
|
|
for target in targets: |
|
|
min_distance = float('inf') |
|
|
for caption in captions: |
|
|
target_box, caption_box = target["box"], caption["box"] |
|
|
distance = abs(target_box[0] - caption_box[0]) + abs(target_box[3] - caption_box[1]) |
|
|
if distance < min_distance: |
|
|
min_distance = distance |
|
|
correct_match = (target, caption) |
|
|
matches.append(correct_match) |
|
|
for target, caption in matches: |
|
|
target_box, caption_box = target["box"], caption["box"] |
|
|
union_box = [ |
|
|
min(target_box[0], caption_box[0]), |
|
|
min(target_box[1], caption_box[1]), |
|
|
max(target_box[2], caption_box[2]), |
|
|
max(target_box[3], caption_box[3]), |
|
|
] |
|
|
figs.append({"box": union_box, "score": 1.0}) |
|
|
return figs |
