Spaces:
Running on Zero
Running on Zero
| import gc | |
| import logging | |
| import os | |
| import re | |
| import subprocess | |
| import sys | |
| import tempfile | |
| from pathlib import Path | |
| from typing import Dict, List | |
| import torch | |
| def _log_runtime_versions(): | |
| try: | |
| import transformers | |
| logger.info("Transformers version: %s", transformers.__version__) | |
| except Exception as exc: | |
| logger.warning("Could not read Transformers version: %s", exc) | |
| logger.info("Torch version: %s | CUDA available: %s", torch.__version__, torch.cuda.is_available()) | |
| logger = logging.getLogger("featurex-janus") | |
| MODEL_ID = os.environ.get("JANUS_MODEL_ID", "ZrH42/Janus-Pro-CXR-Final") | |
| REPO_URL = "https://github.com/ZrH42/Janus-Pro-CXR.git" | |
| REPO_DIR = Path(os.environ.get("JANUS_REPO_DIR", "/tmp/Janus-Pro-CXR")) | |
| JANUS_ENABLED = os.environ.get("JANUS_ENABLED", "true").strip().lower() in {"1", "true", "yes", "on"} | |
| BILINGUAL = os.environ.get("JANUS_BILINGUAL", "false").strip().lower() in {"1", "true", "yes", "on"} | |
| _PROCESSOR = None | |
| _MODEL = None | |
| _INITIALIZED = False | |
| def _clone_repo(): | |
| if (REPO_DIR / "janus").exists(): | |
| return | |
| REPO_DIR.parent.mkdir(parents=True, exist_ok=True) | |
| logger.info("Cloning Janus-Pro-CXR repository to %s", REPO_DIR) | |
| subprocess.run( | |
| ["git", "clone", "--depth", "1", REPO_URL, str(REPO_DIR)], | |
| check=True, | |
| stdout=subprocess.DEVNULL, | |
| ) | |
| def _patch_compatibility(): | |
| siglip = REPO_DIR / "janus" / "models" / "siglip_vit.py" | |
| if siglip.exists(): | |
| text = siglip.read_text(encoding="utf-8") | |
| old = "dpr = [drop_path_rate * i / (depth - 1) for i in torch.linspace(0, depth - 1, depth)]" | |
| if old in text: | |
| text = text.replace(old, "dpr = [drop_path_rate * i / max(depth - 1, 1) for i in range(depth)]") | |
| siglip.write_text(text, encoding="utf-8") | |
| if str(REPO_DIR) not in sys.path: | |
| sys.path.insert(0, str(REPO_DIR)) | |
| try: | |
| from janus.models.modeling_vlm import MultiModalityCausalLM | |
| if not hasattr(MultiModalityCausalLM, "all_tied_weights_keys"): | |
| MultiModalityCausalLM.all_tied_weights_keys = property(lambda self: {}) | |
| except Exception: | |
| pass | |
| def _initialize(): | |
| global _PROCESSOR, _MODEL, _INITIALIZED | |
| if _INITIALIZED: | |
| return | |
| _log_runtime_versions() | |
| if not JANUS_ENABLED: | |
| _INITIALIZED = True | |
| return | |
| if not torch.cuda.is_available(): | |
| raise RuntimeError("Janus requires CUDA/ZeroGPU. The analyze function must execute on a ZeroGPU Space.") | |
| _clone_repo() | |
| _patch_compatibility() | |
| from transformers import AutoModelForCausalLM | |
| from janus.models import VLChatProcessor | |
| logger.info("Loading Janus processor: %s", MODEL_ID) | |
| _PROCESSOR = VLChatProcessor.from_pretrained(MODEL_ID) | |
| logger.info("Loading Janus model: %s", MODEL_ID) | |
| _MODEL = AutoModelForCausalLM.from_pretrained( | |
| MODEL_ID, | |
| trust_remote_code=True, | |
| low_cpu_mem_usage=False, | |
| ).to(torch.bfloat16).cuda().eval() | |
| _INITIALIZED = True | |
| logger.info("Janus model ready on CUDA") | |
| def _prompt(language: str) -> str: | |
| if language == "ar": | |
| return ( | |
| "حلل صورة أشعة الصدر هذه كمساعد في طب الأشعة. اكتب تقريراً طبياً موجزاً باللغة العربية فقط. " | |
| "استخدم العناوين العربية التالية عند الحاجة: النتائج، الانطباع، التوصيات. " | |
| "يجب أن يكون نص التقرير نفسه بالعربية بالكامل، ولا تكتب FINDINGS أو IMPRESSION أو RECOMMENDATIONS، " | |
| "ولا تعِد كتابة التقرير بالإنجليزية. صف فقط العلامات الشعاعية المرئية. " | |
| "لا تخترع أجهزة أو قياسات أو فحوصاً سابقة أو علاجاً أو تاريخاً مرضياً. " | |
| "استخدم صياغة حذرة عند وجود عدم يقين. احتفظ بالأرقام والاختصارات الطبية الضرورية فقط." | |
| ) | |
| return ( | |
| "Analyze this chest X-ray as a radiology assistant. Return a concise report using these headings " | |
| "exactly when applicable: FINDINGS, IMPRESSION, RECOMMENDATIONS. " | |
| "Describe only visible radiographic findings. Do not invent devices, measurements, prior studies, " | |
| "treatment, or patient history. Use cautious wording for uncertain findings. Write the report in English." | |
| ) | |
| def _extract_text(result) -> str: | |
| if isinstance(result, str): | |
| return result.strip() | |
| if isinstance(result, dict): | |
| for key in ("text", "answer", "response", "generated_text"): | |
| if result.get(key): | |
| return str(result[key]).strip() | |
| return str(result).strip() | |
| def _generate_one(image, language: str) -> str: | |
| from janus.utils.io import load_pil_images | |
| # Janus-Pro-CXR's official inference helper accepts an image path. | |
| # Keep the request image isolated in a temporary file and delete it after generation. | |
| suffix = ".png" | |
| tmp = tempfile.NamedTemporaryFile(suffix=suffix, delete=False) | |
| tmp_path = tmp.name | |
| tmp.close() | |
| image.save(tmp_path, format="PNG") | |
| conversation = [ | |
| { | |
| "role": "User", | |
| "content": f"<image_placeholder>\n{_prompt(language)}", | |
| "images": [tmp_path], | |
| }, | |
| {"role": "Assistant", "content": ""}, | |
| ] | |
| pil_images = load_pil_images(conversation) | |
| prepare = _PROCESSOR( | |
| conversations=conversation, | |
| images=pil_images, | |
| force_batchify=True, | |
| ).to(_MODEL.device) | |
| inputs_embeds = _MODEL.prepare_inputs_embeds(**prepare) | |
| outputs = _MODEL.language_model.generate( | |
| inputs_embeds=inputs_embeds, | |
| attention_mask=prepare.attention_mask, | |
| pad_token_id=_PROCESSOR.tokenizer.eos_token_id, | |
| bos_token_id=_PROCESSOR.tokenizer.bos_token_id, | |
| eos_token_id=_PROCESSOR.tokenizer.eos_token_id, | |
| max_new_tokens=512, | |
| do_sample=False, | |
| use_cache=True, | |
| output_scores=True, | |
| return_dict_in_generate=True, | |
| ) | |
| try: | |
| return _extract_text(_PROCESSOR.tokenizer.decode(outputs.sequences[0].cpu().tolist(), skip_special_tokens=True)) | |
| finally: | |
| try: | |
| os.remove(tmp_path) | |
| except OSError: | |
| pass | |
| def _arabic_char_count(text: str) -> int: | |
| return len(re.findall(r"[\u0600-\u06FF]", text or "")) | |
| def _latin_char_count(text: str) -> int: | |
| return len(re.findall(r"[A-Za-z]", text or "")) | |
| def _has_arabic(text: str) -> bool: | |
| # Require actual Arabic content, not a single Arabic punctuation/character. | |
| return _arabic_char_count(text) >= 8 | |
| def _has_english_letters(text: str) -> bool: | |
| return _latin_char_count(text) >= 8 | |
| def _is_predominantly_arabic(text: str) -> bool: | |
| ar = _arabic_char_count(text) | |
| en = _latin_char_count(text) | |
| return ar >= 8 and ar >= max(8, en * 1.25) | |
| def _is_predominantly_english(text: str) -> bool: | |
| ar = _arabic_char_count(text) | |
| en = _latin_char_count(text) | |
| return en >= 8 and en >= max(8, ar * 1.25) | |
| def _generate_arabic_verified(image, attempts: int = 2) -> str: | |
| """Generate Arabic and return an empty string unless the result is genuinely Arabic. | |
| This is important: returning the last English candidate here used to leak English | |
| into ``report_ar_raw`` and then into the Arabic ResultScreen card. | |
| """ | |
| for attempt in range(1, attempts + 1): | |
| candidate = _generate_one(image, "ar") | |
| if _is_predominantly_arabic(candidate): | |
| logger.info("Janus Arabic report verified on attempt %d", attempt) | |
| return candidate.strip() | |
| logger.warning( | |
| "Janus Arabic attempt %d failed language validation (ar=%d, en=%d); retrying", | |
| attempt, _arabic_char_count(candidate), _latin_char_count(candidate), | |
| ) | |
| # Never expose an unverified candidate as Arabic. inference.py will use the | |
| # deterministic Arabic template instead. | |
| return "" | |
| def generate_reports(image, requested_language: str = "ar") -> Dict[str, object]: | |
| """Generate independently generated, language-validated Arabic and English reports.""" | |
| if not JANUS_ENABLED: | |
| return {"enabled": False, "report_ar": "", "report_en": "", "requested_raw_text": ""} | |
| _initialize() | |
| reports = { | |
| "enabled": True, | |
| "report_ar": "", | |
| "report_en": "", | |
| "requested_raw_text": "", | |
| "arabic_verified": False, | |
| "english_verified": False, | |
| } | |
| with torch.inference_mode(): | |
| english_candidate = _generate_one(image, "en") | |
| if _is_predominantly_english(english_candidate): | |
| reports["report_en"] = english_candidate.strip() | |
| reports["english_verified"] = True | |
| else: | |
| logger.warning( | |
| "Janus English output failed language validation (ar=%d, en=%d); using empty raw text", | |
| _arabic_char_count(english_candidate), _latin_char_count(english_candidate), | |
| ) | |
| if BILINGUAL: | |
| reports["report_ar"] = _generate_arabic_verified(image, attempts=2) | |
| reports["arabic_verified"] = _is_predominantly_arabic(reports["report_ar"]) | |
| # requested_raw_text is always language-safe. It is never a fallback from one | |
| # language into the other. | |
| reports["requested_raw_text"] = ( | |
| reports["report_ar"] if requested_language == "ar" and reports["arabic_verified"] | |
| else reports["report_en"] if requested_language == "en" and reports["english_verified"] | |
| else "" | |
| ) | |
| return reports | |
| def parse_radiology_text(text: str) -> Dict[str, object]: | |
| text = (text or "").strip() | |
| if not text: | |
| return {"findings": "", "impression": "", "recommendations": [], "raw_text": ""} | |
| # Normalize common heading variants without trying to reinterpret medical content. | |
| normalized = re.sub(r"\r\n?", "\n", text) | |
| pattern = re.compile( | |
| r"(?:^|[\n\r]|[.\u061b;])\s*" | |
| r"(FINDINGS?|IMPRESSION|CONCLUSION|RECOMMENDATIONS?|RECOMMENDATION|النتائج|الملاحظات|الانطباع|التوصيات)" | |
| r"\s*:?\s*", | |
| flags=re.IGNORECASE, | |
| ) | |
| matches = list(pattern.finditer(normalized)) | |
| sections: Dict[str, str] = {} | |
| for i, match in enumerate(matches): | |
| raw_heading = match.group(1).lower() | |
| start = match.end() | |
| end = matches[i + 1].start() if i + 1 < len(matches) else len(normalized) | |
| body = normalized[start:end].strip(" \n:-•") | |
| if raw_heading.startswith("finding") or raw_heading in {"النتائج", "الملاحظات"}: | |
| sections["findings"] = body | |
| elif raw_heading in {"impression", "conclusion", "الانطباع"}: | |
| sections["impression"] = body | |
| else: | |
| sections["recommendations"] = body | |
| if not sections: | |
| sections["findings"] = normalized | |
| sections["impression"] = normalized | |
| recs = [] | |
| if sections.get("recommendations"): | |
| for line in re.split(r"\n+", sections["recommendations"]): | |
| line = re.sub(r"^\s*(?:[-*•]|\d+[.)])\s*", "", line).strip() | |
| if line: | |
| recs.append(line) | |
| return { | |
| "findings": sections.get("findings", ""), | |
| "impression": sections.get("impression", ""), | |
| "recommendations": recs, | |
| "raw_text": text, | |
| } | |