143 lines
4.4 KiB
Python
143 lines
4.4 KiB
Python
import io, base64, math
|
|
import uuid
|
|
|
|
from PIL import Image, ImageOps, ImageFilter
|
|
from openai import AsyncOpenAI
|
|
from pathlib import Path
|
|
from pdf2image import convert_from_path
|
|
import asyncio
|
|
|
|
# 请求模型地址
|
|
model_base_url = ""
|
|
|
|
async def img_to_base64_qwen3vl(original, path, max_side=2096):
|
|
img = Image.open(path)
|
|
img = ImageOps.exif_transpose(img).convert("RGB")
|
|
|
|
w, h = img.size
|
|
s = min(1.0, max_side / max(w, h))
|
|
if s < 1:
|
|
img = img.resize((int(w * s), int(h * s)), Image.Resampling.LANCZOS)
|
|
|
|
img = img.filter(ImageFilter.UnsharpMask(1, 120, 3))
|
|
|
|
w, h = img.size
|
|
nw, nh = math.ceil(w / 32) * 32, math.ceil(h / 32) * 32
|
|
img = ImageOps.expand(
|
|
img,
|
|
((nw - w) // 2, (nh - h) // 2, nw - w - (nw - w) // 2, nh - h - (nh - h) // 2),
|
|
fill=(128, 128, 128)
|
|
)
|
|
|
|
buf = io.BytesIO()
|
|
img.save(buf, format="JPEG", quality=90, subsampling=0, optimize=True)
|
|
b64 = base64.b64encode(buf.getvalue()).decode()
|
|
img.save(f"./SC/{str(uuid.uuid4())}_{original}", format="JPEG", quality=90, subsampling=0, optimize=True)
|
|
return f"data:image/jpeg;base64,{b64}"
|
|
|
|
|
|
def pdf_to_images_sync(original, pdf_path, dpi=300, fmt="png"):
|
|
try:
|
|
poppler_path = r"C:\Users\FLYF\Downloads\Release-25.12.0-0\poppler-25.12.0\Library\bin"
|
|
|
|
pdf_path = Path(pdf_path)
|
|
id = str(uuid.uuid4())
|
|
|
|
output_dir = Path(f"./TEMP/{pdf_path.stem}_{id}"[:200])
|
|
output_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
images = convert_from_path(
|
|
pdf_path,
|
|
dpi=dpi,
|
|
poppler_path=poppler_path
|
|
)
|
|
|
|
img_paths = []
|
|
ima_names = []
|
|
|
|
for i, img in enumerate(images, start=1):
|
|
out_path = output_dir / f"{pdf_path.stem}_{i}.{fmt}"
|
|
img.save(out_path)
|
|
img_paths.append(str(out_path))
|
|
ima_names.append(f"{original}_{i}.{fmt}")
|
|
|
|
# output = {
|
|
# "pdf_name": original,
|
|
# "result": img_paths,
|
|
# "state": "success"
|
|
# }
|
|
|
|
return img_paths, ima_names
|
|
except Exception as e:
|
|
return img_paths, ima_names
|
|
|
|
|
|
async def pdf_to_images_async(*args, **kwargs):
|
|
return await asyncio.to_thread(pdf_to_images_sync, *args, **kwargs)
|
|
|
|
|
|
with open("SystemPrompt", 'r', encoding='utf-8') as f:
|
|
SystemPrompt = f.read()
|
|
|
|
client = AsyncOpenAI(api_key="1", base_url=model_base_url)
|
|
|
|
|
|
async def qwen3vl(original, img_path, prompt):
|
|
try:
|
|
img = await img_to_base64_qwen3vl(original, img_path)
|
|
response = await client.chat.completions.create(
|
|
model="Qwen3-VL-FLYFAI",
|
|
messages=[
|
|
{"role": "system", "content": SystemPrompt},
|
|
{
|
|
"role": "user",
|
|
"content": [
|
|
{
|
|
"type": "image_url",
|
|
"image_url": {
|
|
"url": img
|
|
}
|
|
},
|
|
{
|
|
"type": "text",
|
|
"text": prompt
|
|
}
|
|
]
|
|
}
|
|
],
|
|
temperature=0.1,
|
|
top_p=0.3,
|
|
max_tokens=4096,
|
|
timeout=90
|
|
)
|
|
output = {
|
|
"image_name": original,
|
|
"result": response.choices[0].message.content,
|
|
"state": "success"
|
|
}
|
|
return output
|
|
except Exception as e:
|
|
return {"image_name": original, "result": str(e), "state": "error"}
|
|
|
|
|
|
async def run_one(semaphore, prompt, original, img_path, retries=3):
|
|
async with semaphore:
|
|
for i in range(retries):
|
|
try:
|
|
return await asyncio.wait_for(qwen3vl(original, img_path, prompt), timeout=90)
|
|
except Exception as e:
|
|
if i == retries - 1:
|
|
return {"image_name": original, "result": str(e), "state": "error"}
|
|
await asyncio.sleep(3)
|
|
|
|
|
|
async def run_images(original, img_paths, prompt, max_concurrency=30):
|
|
sem = asyncio.Semaphore(max_concurrency)
|
|
tasks = [run_one(sem, prompt, original, p) for original, p in zip(original, img_paths)]
|
|
return await asyncio.gather(*tasks)
|
|
|
|
|
|
async def run_pdf(original, pdf_path, prompt):
|
|
img_paths, ima_names = await pdf_to_images_async(original, pdf_path)
|
|
result = await run_images(ima_names, img_paths, prompt)
|
|
return result |