smart-models/Placebo_AI

1import os2import json3import sys4import fitz5import requests6import base647from tqdm import tqdm8from PIL import Image9 10import pytesseract11pytesseract.pytesseract.tesseract_cmd = r'C:\Program Files\Tesseract-OCR\tesseract.exe'12os.environ["PATH"] += os.pathsep + r'C:\Program Files\Tesseract-OCR'13 14OLLAMA_URL = "http://127.0.0.1:11434/api/generate"15 16def check_ollama_status():17    try:18        response = requests.get("http://127.0.0.1:11434/api/tags", timeout=5)19        if response.status_code == 200:20            models = [m.get("name") for m in response.json().get("models", [])]21            print(f"Ollama is running. Available models: {models}")22            return "moondream:latest" in models or "moondream" in models23    except Exception as e:24        print(f"Warning: Could not connect to Ollama at localhost:11434. VLM descriptions will be skipped. ({e})")25    return False26 27def get_vlm_description(image_path):28    prompt = (29        "Describe any anatomical diagrams, histology slides, physiological pathways, clinical charts, "30        "flowcharts, medical algorithms, or dosage tables on this page. If you see a process or cycle, "31        "extract it as a step-by-step logical sequence (Step 1 -> Step 2). Keep descriptions concise and medical-grade. "32        "Do not write long paragraphs. If none, say 'No visual data'."33    )34    try:35        with open(image_path, "rb") as image_file:36            img_str = base64.b64encode(image_file.read()).decode('utf-8')37        38        payload = {39            "model": "moondream",40            "prompt": prompt,41            "images": [img_str],42            "stream": False43        }44        45        response = requests.post(OLLAMA_URL, json=payload, timeout=90)46        if response.status_code == 200:47            return response.json().get("response", "").strip()48        else:49            return f"Ollama VLM Error: {response.status_code}"50    except Exception as e:51        return f"VLM skipped: {e}"52 53def clean_filename(name):54    return name.replace(" ", "_").replace("&", "and").replace("'", "").replace("(", "").replace(")", "").replace("[", "").replace("]", "").replace(",", "")55 56def process_anatomy(pilot=False):57    anatomy_dir = r"d:\sample chatbot\MBBS\books\anatomy"58    master_jsonl_dir = os.path.join(os.path.dirname(os.path.dirname(__file__)), "data", "processed")59    60    has_vlm = check_ollama_status()61    62    pdf_files = []63    for root, dirs, files in os.walk(anatomy_dir):64        for file in files:65            if file.lower().endswith(".pdf") and not file.startswith("._"):66                pdf_files.append({67                    "book": file,68                    "path": os.path.join(root, file)69                })70                71    pdf_files = sorted(pdf_files, key=lambda x: x["book"])72    73    if pilot:74        print("\n=== PILOT MODE ACTIVE ===")75        print("Restricting processing to 1 book, and 2 pages per book.")76        if pdf_files:77            pdf_files = [pdf_files[0]]78            79    jsonl_name = "anatomy_vlm_data.jsonl"80    jsonl_path = os.path.join(master_jsonl_dir, jsonl_name)81    mode = "w" if pilot else "a"82    83    print(f"Total anatomy books to process: {len(pdf_files)}")84    85    for idx, book_info in enumerate(pdf_files):86        book_name = book_info["book"]87        path = book_info["path"]88        89        book_clean = clean_filename(book_name.replace(".pdf", ""))90        subject_clean = "anatomy"91        92        img_out_dir = os.path.join(os.path.dirname(os.path.dirname(__file__)), "data", "mbbs_processed", subject_clean, book_clean)93        os.makedirs(img_out_dir, exist_ok=True)94        95        print(f"\n[{idx+1}/{len(pdf_files)}] Processing book: {book_name}")96        97        try:98            doc = fitz.open(path)99            total_pages = doc.page_count100        except Exception as e:101            print(f"  Error opening PDF {book_name}: {e}")102            continue103            104        processed_count = 0105        106        with open(jsonl_path, mode, encoding="utf-8") as f_out:107            for page_idx in tqdm(range(total_pages), desc=f"Processing {book_clean[:30]}"):108                page_num = page_idx + 1109                110                page = doc[page_idx]111                text = page.get_text().strip()112                text_len = len(text)113                114                if text_len < 50 and len(page.get_images()) == 0:115                    continue116                    117                combined_content = text118                has_images = len(page.get_images()) > 0119                120                img_name = f"page_{page_num:04d}.png"121                img_path = os.path.join(img_out_dir, img_name)122                123                try:124                    pix = page.get_pixmap(dpi=150)125                    pix.save(img_path)126                except Exception as ex:127                    print(f"  Warning: failed to render page image {page_num}: {ex}")128                129                if text_len < 100 and has_images:130                    try:131                        ocr_text = pytesseract.image_to_string(Image.open(img_path))132                        if ocr_text.strip():133                            combined_content = ocr_text134                    except Exception as ocr_err:135                        combined_content = f"OCR Failed: {ocr_err}"136                137                vlm_desc = "No visual data"138                if has_images and has_vlm:139                    vlm_desc = get_vlm_description(img_path)140                141                if vlm_desc and vlm_desc != "No visual data":142                    combined_content += f"\n\nVISUAL DATA:\n{vlm_desc}"143                144                page_data = {145                    "subject": "anatomy",146                    "book_name": book_name,147                    "page_number": page_num,148                    "image_path": os.path.abspath(img_path),149                    "content": combined_content,150                    "category": "anatomy",151                    "track": "mbbs"152                }153                154                f_out.write(json.dumps(page_data, ensure_ascii=False) + "\n")155                f_out.flush()156                157                processed_count += 1158                159                if pilot and processed_count >= 2:160                    break161        162        doc.close()163        print(f"  Successfully processed {processed_count} pages for {book_name}.")164 165if __name__ == "__main__":166    is_pilot = "--full" not in sys.argv167    process_anatomy(pilot=is_pilot)168