diff --git a/config/model_settings.py b/config/model_settings.py index 3128e09..d7eeb46 100644 --- a/config/model_settings.py +++ b/config/model_settings.py @@ -22,8 +22,8 @@ MODEL_PATH = "deepseek-ai/DeepSeek-OCR" # change to your model path # Omnidocbench images path: run_dpsk_ocr_eval_batch.py -INPUT_PATH = "/workspace/input" -OUTPUT_PATH = "/workspace/output" +INPUT_PATH = "/workspace/test/input" +OUTPUT_PATH = "/workspace/test/output" # PROMPT = f"{PROMPT_TEXT.strip()}" PROMPT = "\n<|grounding|>Convert the document to markdown." # PROMPT = '\nFree OCR.' diff --git a/test/__init__.py b/test/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/test/input/2016-08556-211156.pdf b/test/input/2016-08556-211156.pdf new file mode 100644 index 0000000..f552e7b Binary files /dev/null and b/test/input/2016-08556-211156.pdf differ diff --git a/test/output/2016-08556-211156_layouts.pdf b/test/output/2016-08556-211156_layouts.pdf new file mode 100644 index 0000000..803f922 Binary files /dev/null and b/test/output/2016-08556-211156_layouts.pdf differ diff --git a/test/output/images/0_0.jpg b/test/output/images/0_0.jpg new file mode 100644 index 0000000..4cc8b40 Binary files /dev/null and b/test/output/images/0_0.jpg differ diff --git a/test/output/result/2016-08556-211156.json b/test/output/result/2016-08556-211156.json new file mode 100644 index 0000000..73bb350 --- /dev/null +++ b/test/output/result/2016-08556-211156.json @@ -0,0 +1,12 @@ +{ + "filename": "2016-08556-211156.pdf", + "model": { + "ocr_model": "deepseek-ocr" + }, + "time": { + "duration_sec": "6.59", + "started_at": 1762395627.3137395, + "ended_at": 1762395633.9023185 + }, + "parsed": "\n수신자 한국수출입은행장 \n\n참조 EDCF Operations Department 2 \n\n제 목 방글라데시 반다주리 상수도 사업 컨설턴트 고용을 위한 문제유발 행위 불개입 확약서 \n\n1. 귀 은행의 무궁한 발전을 기원합니다. \n\n2. 표제 사업 컨설턴트 고용을 위한 제안요청서 조항에 따라 입찰 참여를 위한 \"문제유발행위 불개입 확약 서\"를 \n\n첨부와 같이 제출하오니, 참조해주시기 바랍니다. \n\n* 첨부: 문제유발행위 불개입 확약서 원본 1부. \n\n주식회사 삼안 대표이사 \n\n![](images/0_0.jpg)\n\n \n\n수신처 : Ms. Jiyoon Park, Sr. Loan Officer \n\n문서번호 201609-4495 (2016-09-22) \n\n서울 광진구 광나루로56실 85 프라임센터 34층 해외사업실 \n\n전화 02)6488-8095 \n\n담당 : 김관영 \n\nFAX 02)6488-8080 \n\n/ http://www.samaneng.com \n\n이메일 shkim5@samaneng.com<｜end▁of▁sentence｜>\n<--- Page Split --->\n" +} \ No newline at end of file diff --git a/test/test.py b/test/test.py new file mode 100644 index 0000000..3afffb8 --- /dev/null +++ b/test/test.py @@ -0,0 +1,397 @@ +import io +import json +import os +import re +import time + +import config.model_settings as config +import fitz +import img2pdf +import numpy as np +from config.env_setup import setup_environment +from PIL import Image, ImageDraw, ImageFont, ImageOps +from services.deepseek_ocr import DeepseekOCRForCausalLM +from services.process.image_process import DeepseekOCRProcessor +from services.process.ngram_norepeat import NoRepeatNGramLogitsProcessor +from tqdm import tqdm +from vllm import LLM, SamplingParams +from vllm.model_executor.models.registry import ModelRegistry + +setup_environment() + +ModelRegistry.register_model("DeepseekOCRForCausalLM", DeepseekOCRForCausalLM) + + +class Colors: + RED = "\033[31m" + GREEN = "\033[32m" + YELLOW = "\033[33m" + BLUE = "\033[34m" + RESET = "\033[0m" + + +# --- PDF/Image Processing Functions (from run_dpsk_ocr_*.py) --- + + +def pdf_to_images_high_quality(pdf_path, dpi=144): + images = [] + pdf_document = fitz.open(pdf_path) + zoom = dpi / 72.0 + matrix = fitz.Matrix(zoom, zoom) + for page_num in range(pdf_document.page_count): + page = pdf_document[page_num] + pixmap = page.get_pixmap(matrix=matrix, alpha=False) + Image.MAX_IMAGE_PIXELS = None + img_data = pixmap.tobytes("png") + img = Image.open(io.BytesIO(img_data)) + if img.mode in ("RGBA", "LA"): + background = Image.new("RGB", img.size, (255, 255, 255)) + background.paste(img, mask=img.split()[-1] if img.mode == "RGBA" else None) + img = background + images.append(img) + pdf_document.close() + return images + + +def pil_to_pdf_img2pdf(pil_images, output_path): + if not pil_images: + return + image_bytes_list = [] + for img in pil_images: + if img.mode != "RGB": + img = img.convert("RGB") + img_buffer = io.BytesIO() + img.save(img_buffer, format="JPEG", quality=95) + image_bytes_list.append(img_buffer.getvalue()) + try: + pdf_bytes = img2pdf.convert(image_bytes_list) + with open(output_path, "wb") as f: + f.write(pdf_bytes) + except Exception as e: + print(f"Error creating PDF: {e}") + + +def re_match(text): + pattern = r"(<\|ref\|>(.*?)<\|/ref\|><\|det\|>(.*?)<\|/det\|>)" + matches = re.findall(pattern, text, re.DOTALL) + mathes_image = [m[0] for m in matches if "<|ref|>image<|/ref|>" in m[0]] + mathes_other = [m[0] for m in matches if "<|ref|>image<|/ref|>" not in m[0]] + return matches, mathes_image, mathes_other + + +def extract_coordinates_and_label(ref_text, image_width, image_height): + try: + label_type = ref_text[1] + cor_list = eval(ref_text[2]) + return (label_type, cor_list) + except Exception as e: + print(f"Error extracting coordinates: {e}") + return None + + +def draw_bounding_boxes(image, refs, jdx=None): + image_width, image_height = image.size + img_draw = image.copy() + draw = ImageDraw.Draw(img_draw) + overlay = Image.new("RGBA", img_draw.size, (0, 0, 0, 0)) + draw2 = ImageDraw.Draw(overlay) + font = ImageFont.load_default() + img_idx = 0 + + for i, ref in enumerate(refs): + result = extract_coordinates_and_label(ref, image_width, image_height) + if not result: + continue + + label_type, points_list = result + color = ( + np.random.randint(0, 200), + np.random.randint(0, 200), + np.random.randint(0, 255), + ) + color_a = color + (20,) + + for points in points_list: + x1, y1, x2, y2 = [ + int(p / 999 * (image_width if i % 2 == 0 else image_height)) + for i, p in enumerate(points) + ] + + if label_type == "image": + try: + cropped = image.crop((x1, y1, x2, y2)) + img_filename = ( + f"{jdx}_{img_idx}.jpg" if jdx is not None else f"{img_idx}.jpg" + ) + cropped.save( + os.path.join(config.OUTPUT_PATH, "images", img_filename) + ) + img_idx += 1 + except Exception as e: + print(f"Error cropping image: {e}") + + width = 4 if label_type == "title" else 2 + draw.rectangle([x1, y1, x2, y2], outline=color, width=width) + draw2.rectangle( + [x1, y1, x2, y2], fill=color_a, outline=(0, 0, 0, 0), width=1 + ) + + text_x, text_y = x1, max(0, y1 - 15) + text_bbox = draw.textbbox((0, 0), label_type, font=font) + text_width, text_height = ( + text_bbox[2] - text_bbox[0], + text_bbox[3] - text_bbox[1], + ) + draw.rectangle( + [text_x, text_y, text_x + text_width, text_y + text_height], + fill=(255, 255, 255, 30), + ) + draw.text((text_x, text_y), label_type, font=font, fill=color) + + img_draw.paste(overlay, (0, 0), overlay) + return img_draw + + +def process_image_with_refs(image, ref_texts, jdx=None): + return draw_bounding_boxes(image, ref_texts, jdx) + + +def load_image(image_path): + try: + image = Image.open(image_path).convert("RGB") + return ImageOps.exif_transpose(image) + except Exception as e: + print(f"Error loading image {image_path}: {e}") + return None + + +# --- Main OCR Processing Logic --- + + +def process_pdf(llm, sampling_params, pdf_path): + print(f"{Colors.GREEN}Processing PDF: {pdf_path}{Colors.RESET}") + base_name = os.path.basename(pdf_path) + file_name_without_ext = os.path.splitext(base_name)[0] + + images = pdf_to_images_high_quality(pdf_path) + if not images: + print( + f"{Colors.YELLOW}Could not extract images from {pdf_path}. Skipping.{Colors.RESET}" + ) + return + + batch_inputs = [] + processor = DeepseekOCRProcessor() + for image in tqdm(images, desc="Pre-processing PDF pages"): + batch_inputs.append( + { + "prompt": config.PROMPT, + "multi_modal_data": { + "image": processor.tokenize_with_images( + images=[image], bos=True, eos=True, cropping=config.CROP_MODE + ) + }, + } + ) + + start_time = time.time() + outputs_list = llm.generate(batch_inputs, sampling_params=sampling_params) + end_time = time.time() + + contents_det = "" + contents = "" + draw_images = [] + for i, (output, img) in enumerate(zip(outputs_list, images)): + content = output.outputs[0].text + if "<｜end of sentence｜>" in content: + content = content.replace("<｜end of sentence｜>", "") + elif config.SKIP_REPEAT: + continue + + page_num_separator = "\n<--- Page Split --->\n" + contents_det += content + page_num_separator + + matches_ref, matches_images, mathes_other = re_match(content) + result_image = process_image_with_refs(img.copy(), matches_ref, jdx=i) + draw_images.append(result_image) + + for idx, match in enumerate(matches_images): + content = content.replace(match, f"![](images/{i}_{idx}.jpg)\n") + for match in mathes_other: + content = ( + content.replace(match, "") + .replace("\\coloneqq", ":=") + .replace("\\eqqcolon", "=:") + .replace("\n\n\n", "\n\n") + ) + + contents += content + page_num_separator + + # Save results + json_path = os.path.join( + f"{config.OUTPUT_PATH}/result", f"{file_name_without_ext}.json" + ) + pdf_out_path = os.path.join( + config.OUTPUT_PATH, f"{file_name_without_ext}_layouts.pdf" + ) + + duration = end_time - start_time + output_data = { + "filename": base_name, + "model": {"ocr_model": "deepseek-ocr"}, + "time": { + "duration_sec": f"{duration:.2f}", + "started_at": start_time, + "ended_at": end_time, + }, + "parsed": contents, + } + + with open(json_path, "w", encoding="utf-8") as f: + json.dump(output_data, f, ensure_ascii=False, indent=4) + + pil_to_pdf_img2pdf(draw_images, pdf_out_path) + print( + f"{Colors.GREEN}Finished processing {pdf_path}. Results saved in {config.OUTPUT_PATH}{Colors.RESET}" + ) + + +def process_image(llm, sampling_params, image_path): + print(f"{Colors.GREEN}Processing Image: {image_path}{Colors.RESET}") + base_name = os.path.basename(image_path) + file_name_without_ext = os.path.splitext(base_name)[0] + + image = load_image(image_path) + if image is None: + return + + processor = DeepseekOCRProcessor() + image_features = processor.tokenize_with_images( + images=[image], bos=True, eos=True, cropping=config.CROP_MODE + ) + + request = { + "prompt": config.PROMPT, + "multi_modal_data": {"image": image_features}, + } + + start_time = time.time() + outputs = llm.generate([request], sampling_params) + end_time = time.time() + result_out = outputs[0].outputs[0].text + + print(result_out) + + # Save results + output_json_path = os.path.join( + f"{config.OUTPUT_PATH}/result", f"{file_name_without_ext}.json" + ) + result_image_path = os.path.join( + config.OUTPUT_PATH, f"{file_name_without_ext}_result_with_boxes.jpg" + ) + + matches_ref, matches_images, mathes_other = re_match(result_out) + result_image = process_image_with_refs(image.copy(), matches_ref) + + processed_text = result_out + for idx, match in enumerate(matches_images): + processed_text = processed_text.replace(match, f"![](images/{idx}.jpg)\n") + for match in mathes_other: + processed_text = ( + processed_text.replace(match, "") + .replace("\\coloneqq", ":=") + .replace("\\eqqcolon", "=:") + .replace("\n\n\n", "\n\n") + ) + + duration = end_time - start_time + output_data = { + "filename": base_name, + "model": {"ocr_model": "deepseek-ocr"}, + "time": { + "duration_sec": f"{duration:.2f}", + "started_at": start_time, + "ended_at": end_time, + }, + "parsed": processed_text, + } + + with open(output_json_path, "w", encoding="utf-8") as f: + json.dump(output_data, f, ensure_ascii=False, indent=4) + + result_image.save(result_image_path) + print( + f"{Colors.GREEN}Finished processing {image_path}. Results saved in {config.OUTPUT_PATH}{Colors.RESET}" + ) + + +def main(): + # --- Model Initialization --- + print(f"{Colors.BLUE}Initializing model...{Colors.RESET}") + llm = LLM( + model=config.MODEL_PATH, + hf_overrides={"architectures": ["DeepseekOCRForCausalLM"]}, + block_size=256, + enforce_eager=False, + trust_remote_code=True, + max_model_len=8192, + swap_space=0, + max_num_seqs=config.MAX_CONCURRENCY, + tensor_parallel_size=1, + gpu_memory_utilization=0.9, + disable_mm_preprocessor_cache=True, + ) + + logits_processors = [ + NoRepeatNGramLogitsProcessor( + ngram_size=20, window_size=50, whitelist_token_ids={128821, 128822} + ) + ] + + sampling_params = SamplingParams( + temperature=0.0, + max_tokens=8192, + logits_processors=logits_processors, + skip_special_tokens=False, + include_stop_str_in_output=True, + ) + print(f"{Colors.BLUE}Model initialized successfully.{Colors.RESET}") + + # --- File Processing --- + input_dir = config.INPUT_PATH + output_dir = config.OUTPUT_PATH + os.makedirs(output_dir, exist_ok=True) + os.makedirs(os.path.join(output_dir, "images"), exist_ok=True) + os.makedirs(os.path.join(output_dir, "result"), exist_ok=True) + + if not os.path.isdir(input_dir): + print( + f"{Colors.RED}Error: Input directory not found at '{input_dir}'{Colors.RESET}" + ) + return + + print(f"Scanning for files in '{input_dir}'...") + for filename in sorted(os.listdir(input_dir)): + input_path = os.path.join(input_dir, filename) + if not os.path.isfile(input_path): + continue + + file_extension = os.path.splitext(filename)[1].lower() + + try: + if file_extension == ".pdf": + process_pdf(llm, sampling_params, input_path) + elif file_extension in [".jpg", ".jpeg", ".png", ".bmp", ".gif", ".webp"]: + process_image(llm, sampling_params, input_path) + else: + print( + f"{Colors.YELLOW}Skipping unsupported file type: {filename}{Colors.RESET}" + ) + except Exception as e: + print( + f"{Colors.RED}An error occurred while processing {filename}: {e}{Colors.RESET}" + ) + + +if __name__ == "__main__": + main()