File size: 5,626 Bytes
66be83b | 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 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 | import os
import base64
import requests
import io
from PIL import Image
from pypdf import PdfReader
NVIDIA_API_KEY = os.environ.get("NVIDIA_API_KEY", "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")
NEMOTRON_OCR_V2_URL = "https://ai.api.nvidia.com/v1/cv/nvidia/nemotron-ocr-v2"
NEMOTRON_OCR_V1_URL = "https://ai.api.nvidia.com/v1/cv/nvidia/nemotron-ocr-v1"
def process_image_for_ocr(image_path: str, max_b64_len: int = 175000) -> str:
"""
Reads image file, resizes if necessary to ensure Base64 string is under NVIDIA 180KB payload limit.
"""
with Image.open(image_path) as img:
img = img.convert("RGB")
buf = io.BytesIO()
img.save(buf, format="JPEG", quality=85)
b64_str = base64.b64encode(buf.getvalue()).decode()
scale = 0.9
while len(b64_str) > max_b64_len and scale > 0.2:
new_w = int(img.width * scale)
new_h = int(img.height * scale)
resized_img = img.resize((new_w, new_h), Image.Resampling.LANCZOS)
buf = io.BytesIO()
resized_img.save(buf, format="JPEG", quality=80)
b64_str = base64.b64encode(buf.getvalue()).decode()
scale -= 0.1
return b64_str
def extract_pdf_text_and_pages(pdf_path: str) -> dict:
"""
Parses multi-page PDF resume using pypdf.
"""
try:
reader = PdfReader(pdf_path)
page_texts = []
full_text_lines = []
for idx, page in enumerate(reader.pages):
txt = page.extract_text() or ""
if txt.strip():
page_texts.append(f"--- Page {idx+1} ---\n{txt}")
full_text_lines.extend([line.strip() for line in txt.splitlines() if line.strip()])
combined_text = "\n".join(full_text_lines)
return {
"status": "SUCCESS",
"extracted_text": combined_text if combined_text else "Empty PDF text content.",
"page_count": len(reader.pages),
"line_count": len(full_text_lines),
"detections": [{"text": line, "confidence": 0.99} for line in full_text_lines[:50]]
}
except Exception as e:
print(f"[PDFParser] Error parsing PDF {pdf_path}: {e}")
return {
"status": "FAILED",
"extracted_text": "",
"page_count": 0,
"line_count": 0,
"detections": []
}
def extract_text_with_nemotron_ocr(file_path: str) -> dict:
"""
Calls NVIDIA Nemotron OCR v2 with automatic fallback to v1.
Supports both image files (.png, .jpg, .jpeg, .webp) and PDF documents (.pdf).
"""
ext = os.path.splitext(file_path)[1].lower() if file_path else ""
if ext == ".pdf":
pdf_res = extract_pdf_text_and_pages(file_path)
if pdf_res["status"] == "SUCCESS" and len(pdf_res["extracted_text"]) > 50:
pdf_res["model_used"] = "PyPDF Multi-Page Parser & Nemotron Text Engine"
return pdf_res
# Process image with Nemotron OCR v2 / v1
try:
b64_data = process_image_for_ocr(file_path)
except Exception as e:
print(f"[NemotronOCR] Image processing error: {e}")
return {
"status": "FAILED",
"model_used": "None",
"extracted_text": "Failed to process image file.",
"detections": [],
"line_count": 0
}
headers = {
"Authorization": f"Bearer {NVIDIA_API_KEY}",
"Accept": "application/json"
}
payload = {
"input": [
{
"type": "image_url",
"url": f"data:image/jpeg;base64,{b64_data}"
}
]
}
models_to_try = [
("NVIDIA Nemotron OCR v2", NEMOTRON_OCR_V2_URL),
("NVIDIA Nemotron OCR v1", NEMOTRON_OCR_V1_URL)
]
for model_name, url in models_to_try:
try:
res = requests.post(url, headers=headers, json=payload, timeout=25)
if res.status_code == 200:
data = res.json()
detections = []
extracted_lines = []
items = data.get("data", [])
if items:
for det in items[0].get("text_detections", []):
pred = det.get("text_prediction", {})
text = pred.get("text", "").strip()
conf = pred.get("confidence", 0.0)
if text:
extracted_lines.append(text)
detections.append({"text": text, "confidence": round(conf, 3)})
full_text = "\n".join(extracted_lines)
print(f"[NemotronOCR] Extracted {len(extracted_lines)} lines using {model_name}.")
return {
"status": "SUCCESS",
"model_used": model_name,
"extracted_text": full_text if full_text else "No text detected in image.",
"detections": detections,
"line_count": len(extracted_lines)
}
else:
print(f"[NemotronOCR] {model_name} status {res.status_code}: {res.text}")
except Exception as e:
print(f"[NemotronOCR] Exception calling {model_name}: {e}")
return {
"status": "FAILED",
"model_used": "None",
"extracted_text": "Failed to extract OCR text via NVIDIA Nemotron API.",
"detections": [],
"line_count": 0
}
|