--- language: - en license: apache-2.0 tags: - medical - radiology - image-captioning - resnet-50 - transformers - pytorch - vision-and-language datasets: - eltorio/ROCOv2-radiology --- # Model Card: ResNet-50 Graph Radiology Captioner ## Model Details - **Model Name:** ResNet-50 Graph Radiology Captioner - **Architecture:** Custom Encoder-Decoder with Graph Propagation - **Language:** English - **Dataset:** eltorio/ROCOv2-radiology - **License:** Apache-2.0 ## Model Description This model is designed to generate automated textual descriptions (captions) for radiology images. It utilizes a hybrid architecture combining Convolutional Neural Networks (CNNs) with Graph Propagation and Transformer decoders. The architecture consists of three main components: 1. **Visual Encoder:** A pre-trained `ResNet-50` (from Hugging Face) extracts visual features from the input image. 2. **Graph Propagation Layer:** A custom layer processes the spatial feature map using a sliding window mechanism (window size $3$) to aggregate neighborhood information over $3$ steps. This allows the model to better capture spatial relationships between anatomical structures. 3. **Textual Decoder:** A standard Transformer Decoder (4 layers, 8 heads) generates the caption token by token, attending to the enhanced visual features. ## Intended Uses - **Primary Use:** Generating preliminary descriptive reports for medical imaging datasets. - **Research:** Investigating graph-based spatial reasoning in medical image captioning. ### Limitations - The model may generate repetitive phrases or hallucinate findings not present in the image. - The inference logic relies on a greedy decoder; beam search could improve results. - Performance is strictly tied to the ROCOv2 dataset distribution; generalizability to other clinical domains may vary. ## Training Data The model was trained on the **ROCOv2-radiology** dataset. - **Dataset Split:** - Train: 56,963 images - Validation: 2,999 images - **Image Modality:** Radiology (X-ray, CT, MRI, etc.) - **Preprocessing:** Images resized to $224 \times 224$, normalized using ImageNet statistics, and augmented with random horizontal flips and rotations ($\pm 10^\circ$). ## Training Procedure The model was trained for 15 epochs on a TPU environment using the CrossEntropyLoss function. ### Hyperparameters | Parameter | Value | | :--- | :--- | | **Batch Size** | 200 | | **Learning Rate** | $1 \times 10^{-4}$ | | **Optimizer** | AdamW | | **Epochs** | 15 | | **Max Sequence Length** | 80 | | **Embedding Dimension** | 512 | | **Feed Forward Dimension** | 1024 | | **Decoder Layers** | 4 | | **Attention Heads** | 8 | | **Graph Window Size** | 3 | | **Graph Propagation Steps** | 3 | ### Training Results The model achieved a steady decrease in training loss and perplexity over 15 epochs. - **Final Training Loss:** $2.4942$ - **Final Perplexity:** $12.1121$ | Epoch | Training Loss | Perplexity | | :--- | :--- | :--- | | 1 | 6.0415 | 420.53 | | 5 | 3.3029 | 27.19 | | 10 | 2.7978 | 16.41 | | 15 | 2.4942 | 12.11 | ## Evaluation ### Inference Sample Below is a comparison between the ground truth caption and the model's generation on a validation set image. **Ground Truth:** "Figure 1: PICC outlined by the arrows." **Generated Caption:** "radiograph of the abdomen showing the presence of a large, well - defined, oval, calcified, and calcifications in the distal part of the right lower quadrant ( arrows )." ## How to Use To use this model, you need to install `torch`, `transformers`, and `pillow`. The model architecture must be defined in your environment to load the weights. ```python import torch import torch.nn as nn from transformers import BertTokenizer, AutoModel from PIL import Image from torchvision import transforms # 1. Define the Model Architecture (Must match training script) class GraphPropagationLayer(nn.Module): def __init__(self, embed_dim, window_size=3, max_steps=3): super().__init__() self.window_size = window_size self.max_steps = max_steps self.msg_linear = nn.Linear(embed_dim, embed_dim) self.activation = nn.Hardtanh() def forward(self, x): B, L, D = x.shape pad = self.window_size // 2 for _ in range(self.max_steps): neighbors = [] for i in range(-pad, pad + 1): neighbors.append(torch.roll(x, shifts=i, dims=1)) stacked = torch.stack(neighbors, dim=2) aggregated = stacked.mean(dim=2) messages = self.msg_linear(aggregated) x = self.activation(x + messages) return x class ResNetGraphCaptioner(nn.Module): def __init__(self, vocab_size, embed_dim, num_layers, num_heads): super().__init__() self.encoder = AutoModel.from_pretrained("microsoft/resnet-50", trust_remote_code=True) self.encoder.config.pooling = None encoder_dim = self.encoder.config.hidden_sizes[-1] self.vis_proj = nn.Linear(encoder_dim, embed_dim) self.graph_layer = GraphPropagationLayer(embed_dim, window_size=3, max_steps=3) self.embedding = nn.Embedding(vocab_size, embed_dim) self.pos_encoder = nn.Parameter(torch.randn(1, 2000, embed_dim)) decoder_layer = nn.TransformerDecoderLayer(d_model=embed_dim, nhead=num_heads, dim_feedforward=1024, batch_first=True) self.decoder = nn.TransformerDecoder(decoder_layer, num_layers=num_layers) self.fc_out = nn.Linear(embed_dim, vocab_size) def forward(self, images, captions): # Encode outputs = self.encoder(pixel_values=images) B, C, H, W = outputs.last_hidden_state.shape visual_features = outputs.last_hidden_state.permute(0, 2, 3, 1).reshape(B, H*W, C) visual_tokens = self.vis_proj(visual_features) memory = self.graph_layer(visual_tokens) # Decode caption_emb = self.embedding(captions) seq_len = caption_emb.size(1) caption_emb = caption_emb + self.pos_encoder[:, :seq_len, :] tgt_mask = torch.triu(torch.ones((seq_len, seq_len)) * float('-inf'), diagonal=1).to(images.device) output = self.decoder(tgt=caption_emb, memory=memory, tgt_mask=tgt_mask) return self.fc_out(output) # 2. Load Checkpoint device = torch.device("cuda" if torch.cuda.is_available() else "cpu") tokenizer = BertTokenizer.from_pretrained('bert-base-uncased') # Initialize model with specific config used during training model = ResNetGraphCaptioner( vocab_size=tokenizer.vocab_size, embed_dim=512, num_layers=4, num_heads=8 ) # Load weights (assuming you have the checkpoint file locally) checkpoint_path = "erfanasghariyan/resnet50_roco_radiocaptioner_big/roco_best_model.bin" model.load_state_dict(torch.load(checkpoint_path, map_location=device)) model.to(device) model.eval() # 3. Inference Function def generate_caption(image_path, model, tokenizer, max_len=80): transform = transforms.Compose([ transforms.Resize((224, 224)), transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), ]) image = Image.open(image_path).convert('RGB') img_tensor = transform(image).unsqueeze(0).to(device) input_ids = [tokenizer.cls_token_id] with torch.no_grad(): for _ in range(max_len): tensor_input = torch.tensor([input_ids], device=device) outputs = model(img_tensor, tensor_input) next_token = outputs[0, -1, :].argmax().item() if next_token == tokenizer.sep_token_id: break input_ids.append(next_token) return tokenizer.decode(input_ids[1:], skip_special_tokens=True) # Example Usage # caption = generate_caption("path/to/xray.jpg", model, tokenizer) # print(caption)