Spaces:
Sleeping
Sleeping
| import os | |
| import logging | |
| import numpy as np | |
| import pandas as pd | |
| from huggingface_hub import snapshot_download | |
| # Configure logging | |
| logger = logging.getLogger(__name__) | |
| class PrecomputedModel: | |
| def __init__(self, data_dir="data"): | |
| self.data_dir = data_dir | |
| self.image_embeddings = None | |
| self.text_embeddings = None | |
| self.labels = None | |
| self._load_data() | |
| def _load_data(self): | |
| """Loads precomputed embeddings and labels.""" | |
| img_emb_path = os.path.join(self.data_dir, "precomputed_image_embeddings.npz") | |
| txt_emb_path = os.path.join(self.data_dir, "precomputed_text_embeddings.npz") | |
| labels_path = os.path.join(self.data_dir, "cxr14_subset_labels.csv") | |
| # Text embeddings are strictly required for Zero-Shot | |
| if not os.path.exists(txt_emb_path): | |
| raise FileNotFoundError(f"Missing required text embeddings: {txt_emb_path}") | |
| logger.info("Loading precomputed text embeddings...") | |
| with np.load(txt_emb_path) as data: | |
| self.text_embeddings = {key: data[key] for key in data} | |
| # Image embeddings (Optional, only for benchmarking) | |
| if os.path.exists(img_emb_path): | |
| logger.info("Loading precomputed image embeddings...") | |
| with np.load(img_emb_path) as data: | |
| self.image_embeddings = {key: data[key] for key in data} | |
| else: | |
| logger.warning("Precomputed image embeddings not found. Benchmarking features will be disabled.") | |
| # Labels (Optional) | |
| if os.path.exists(labels_path): | |
| logger.info("Loading labels...") | |
| self.labels = pd.read_csv(labels_path) | |
| else: | |
| logger.warning("Labels file not found.") | |
| def get_diagnosis_embeddings(self, pos_txt, neg_txt): | |
| """Retrieves embeddings for positive and negative text queries.""" | |
| if pos_txt not in self.text_embeddings: | |
| raise ValueError(f"Positive query '{pos_txt}' not found in precomputed embeddings.") | |
| if neg_txt not in self.text_embeddings: | |
| raise ValueError(f"Negative query '{neg_txt}' not found in precomputed embeddings.") | |
| return self.text_embeddings[pos_txt], self.text_embeddings[neg_txt] | |
| def compute_scores(self, image_ids, pos_emb, neg_emb): | |
| """Computes zero-shot scores for a list of image IDs.""" | |
| scores = [] | |
| valid_ids = [] | |
| for img_id in image_ids: | |
| if img_id not in self.image_embeddings: | |
| continue | |
| img_emb = self.image_embeddings[img_id] | |
| score = self.zero_shot(img_emb, pos_emb, neg_emb) | |
| scores.append(score) | |
| valid_ids.append(img_id) | |
| return valid_ids, scores | |
| def compute_image_text_similarity(image_emb, txt_emb): | |
| """Computes cosine similarity between image and text embeddings.""" | |
| # Image embedding shape: (1, 32, 128) or (32, 128) flattened? | |
| # The notebook says: image_emb = np.reshape(image_emb, (32, 128)) | |
| image_emb = np.reshape(image_emb, (32, 128)) | |
| similarities = [] | |
| for i in range(32): | |
| # cosine similarity | |
| similarity = np.dot(image_emb[i], txt_emb) / (np.linalg.norm(image_emb[i]) * np.linalg.norm(txt_emb)) | |
| similarities.append(similarity) | |
| return np.max(similarities) | |
| def zero_shot(cls, image_emb, pos_txt_emb, neg_txt_emb): | |
| """Computes the zero-shot score (pos_sim - neg_sim).""" | |
| pos_cosine = cls.compute_image_text_similarity(image_emb, pos_txt_emb) | |
| neg_cosine = cls.compute_image_text_similarity(image_emb, neg_txt_emb) | |
| return pos_cosine - neg_cosine | |
| class RawImageModel: | |
| def __init__(self): | |
| self.elixrc_model = None | |
| self.qformer_model = None | |
| self._load_model() | |
| def _load_model(self): | |
| """Loads the TensorFlow model from Hugging Face.""" | |
| try: | |
| import tensorflow as tf | |
| import tensorflow_text as text # Registers the ops | |
| except ImportError: | |
| raise ImportError("TensorFlow or tensorflow-text is not installed. Use precomputed mode or install them.") | |
| logger.info("Checking for GPU acceleration...") | |
| gpus = tf.config.list_physical_devices('GPU') | |
| if gpus: | |
| logger.info(f"Running on GPU: {gpus}") | |
| else: | |
| logger.info("Running on CPU. Expect slower inference.") | |
| logger.info("Downloading model weights from Hugging Face...") | |
| model_path = snapshot_download( | |
| repo_id="google/cxr-foundation", | |
| allow_patterns=['elixr-c-v2-pooled/*', 'pax-elixr-b-text/*'] | |
| ) | |
| logger.info("Loading ELIXR-C (Image Encoder)...") | |
| self.elixrc_model = tf.saved_model.load(os.path.join(model_path, 'elixr-c-v2-pooled')) | |
| logger.info("Loading QFormer (Adapter)...") | |
| self.qformer_model = tf.saved_model.load(os.path.join(model_path, 'pax-elixr-b-text')) | |
| def compute_embeddings(self, image_path): | |
| """Generates embeddings for a raw image file.""" | |
| import tensorflow as tf | |
| import png # pypng | |
| # Load and preprocess image | |
| # This follows the notebook's png_to_tfexample logic but simplified or imported | |
| # For simplicity, implementing the preprocess logic here | |
| try: | |
| # Read image using pypng logic or similar | |
| # Note: The notebook uses pypng to write to BytesIO then TF reads it. | |
| # We can just read the file directly if it's a PNG. | |
| with open(image_path, 'rb') as f: | |
| image_bytes = f.read() | |
| # Create TF Example | |
| example = tf.train.Example() | |
| features = example.features.feature | |
| features['image/encoded'].bytes_list.value.append(image_bytes) | |
| features['image/format'].bytes_list.value.append(b'png') | |
| serialized_example = example.SerializeToString() | |
| # Step 1: ELIXR-C | |
| elixrc_infer = self.elixrc_model.signatures['serving_default'] | |
| elixrc_output = elixrc_infer(input_example=tf.constant([serialized_example])) | |
| elixrc_embedding = elixrc_output['feature_maps_0'].numpy() # Shape (1, 8, 8, 1376) | |
| # Step 2: QFormer | |
| # Initialize text inputs with zeros (as we only want image embeddings) | |
| qformer_input = { | |
| 'image_feature': elixrc_embedding.tolist(), | |
| 'ids': np.zeros((1, 1, 128), dtype=np.int32).tolist(), | |
| 'paddings': np.zeros((1, 1, 128), dtype=np.float32).tolist(), | |
| } | |
| qformer_output = self.qformer_model.signatures['serving_default'](**qformer_input) | |
| elixrb_embeddings = qformer_output['all_contrastive_img_emb'].numpy() # Shape (1, 32, 128) | |
| return elixrb_embeddings | |
| except Exception as e: | |
| logger.error(f"Error computing raw embeddings: {e}") | |
| raise | |