smart-models/Placebo_AI

1import os2import json3import fitz4from PIL import Image5import pytesseract6from concurrent.futures import ProcessPoolExecutor, as_completed7from tqdm import tqdm8 9pytesseract.pytesseract.tesseract_cmd = r'C:\Program Files\Tesseract-OCR\tesseract.exe'10 11def clean_filename(name):12    return name.replace(" ", "_").replace("&", "and").replace("'", "").replace("(", "").replace(")", "").replace("[", "").replace("]", "").replace(",", "")13 14def process_page(args):15    pdf_path, page_num, book_name = args16    try:17        doc = fitz.open(pdf_path)18        page = doc[page_num]19        text = page.get_text().strip()20        21        has_images = len(page.get_images()) > 022        23        book_clean = clean_filename(book_name.replace(".pdf", ""))24        subject_clean = "anatomy"25        26        img_out_dir = os.path.join(os.path.dirname(os.path.dirname(__file__)), "data", "mbbs_processed", subject_clean, book_clean)27        # Create directories safely28        os.makedirs(img_out_dir, exist_ok=True)29        img_name = f"page_{page_num+1:04d}.png"30        img_path = os.path.join(img_out_dir, img_name)31        32        # Save image for UI Gallery33        pix = page.get_pixmap(dpi=150)34        pix.save(img_path)35        36        if len(text) < 100 and has_images:37            # Use OCR if not enough text38            img = Image.open(img_path)39            ocr_text = pytesseract.image_to_string(img)40            if ocr_text.strip():41                text = ocr_text42                43        doc.close()44        45        if len(text.strip()) < 50:46            return None47            48        return {49            "book_name": book_name,50            "page_number": page_num + 1,51            "image_path": os.path.abspath(img_path),52            "content": text.strip(),53            "category": "anatomy",54            "track": "mbbs",55            "subject": "anatomy"56        }57    except Exception as e:58        return None59 60def main():61    anatomy_dir = r"d:\sample chatbot\MBBS\books\anatomy"62    output_jsonl = r"d:\sample chatbot\anatomy_data.jsonl"63    64    if os.path.exists(output_jsonl):65        print(f"Removing existing {output_jsonl} to start fresh...")66        os.remove(output_jsonl)67        68    tasks = []69    print("Gathering PDF info...")70    for root, dirs, files in os.walk(anatomy_dir):71        for filename in files:72            if not filename.lower().endswith(".pdf"):73                continue74            if filename.startswith("._"):75                continue76            77            pdf_path = os.path.join(root, filename)78            try:79                doc = fitz.open(pdf_path)80                total_pages = doc.page_count81                doc.close()82                for i in range(total_pages):83                    tasks.append((pdf_path, i, filename))84            except Exception as e:85                print(f"Failed to open {filename}: {e}")86                87    print(f"Total pages to process: {len(tasks)}")88    89    with open(output_jsonl, 'a', encoding='utf-8') as f_out:90        with ProcessPoolExecutor(max_workers=os.cpu_count() or 4) as executor:91            futures = [executor.submit(process_page, task) for task in tasks]92            93            for future in tqdm(as_completed(futures), total=len(futures), desc="Processing pages"):94                res = future.result()95                if res:96                    f_out.write(json.dumps(res, ensure_ascii=False) + "\n")97                    f_out.flush()98                    99    print(f"Finished! Data saved to {output_jsonl}")100 101if __name__ == "__main__":102    main()103