Update app.py

app.py

@@ -1,125 +1,41 @@
-# app.py
 import gradio as gr
-import torch
+from transformers import BlipProcessor, BlipForConditionalGeneration
 from PIL import Image
-import json
-import os
-from tokenizers import ByteLevelBPETokenizer # Changed from Tokenizer
-from torchvision import transforms
+import requests
 
-# ============================================================
-# 3. Model (Tiny CNN + Transformer Decoder) - Re-define model classes
-# ============================================================
-class CNNEncoder(torch.nn.Module):
-    def __init__(self, embed_dim=128):
-        super().__init__()
-        self.cnn = torch.nn.Sequential(
-            torch.nn.Conv2d(3, 32, 3, 2, 1), torch.nn.ReLU(),
-            torch.nn.Conv2d(32, 64, 3, 2, 1), torch.nn.ReLU(),
-            torch.nn.Conv2d(64, 128, 3, 2, 1), torch.nn.ReLU(),
-            torch.nn.AdaptiveAvgPool2d((1,1))
-        )
-        self.fc = torch.nn.Linear(128, embed_dim)
-    def forward(self, x):
-        x = self.cnn(x)
-        x = x.view(x.size(0), -1)
-        return self.fc(x)
+# Load model and processor
+model_name = "hackergeek/radiology-image-captioning"
+processor = BlipProcessor.from_pretrained(model_name)
+model = BlipForConditionalGeneration.from_pretrained(model_name)
 
-class TransformerDecoder(torch.nn.Module):
-    def __init__(self, vocab_size, embed_dim=128, nhead=4, num_layers=2, max_len=40):
-        super().__init__()
-        self.embed = torch.nn.Embedding(vocab_size, embed_dim)
-        decoder_layer = torch.nn.TransformerDecoderLayer(d_model=embed_dim, nhead=nhead)
-        self.decoder = torch.nn.TransformerDecoder(decoder_layer, num_layers=num_layers)
-        self.fc_out = torch.nn.Linear(embed_dim, vocab_size)
-        self.pos_embed = torch.nn.Embedding(max_len, embed_dim)
-
-    def forward(self, tgt, memory):
-        positions = torch.arange(0, tgt.shape[1], device=tgt.device).unsqueeze(0)
-        tgt_emb = self.embed(tgt) + self.pos_embed(positions)
-        memory = memory.unsqueeze(0)
-        out = self.decoder(tgt_emb.transpose(0,1), memory)
-        return self.fc_out(out.transpose(0,1))
-
-class ImageCaptionModel(torch.nn.Module):
-    def __init__(self, vocab_size, embed_dim=128):
-        super().__init__()
-        self.encoder = CNNEncoder(embed_dim)
-        self.decoder = TransformerDecoder(vocab_size, embed_dim)
-    def forward(self, images, captions):
-        feats = self.encoder(images)
-        return self.decoder(captions, feats)
-
-# ============================================================
-# Load the tokenizer and model manually
-# ============================================================
-
-# Load config
-with open("hackergeek/radiology-image-captioning/config.json", "r") as f:
-    config = json.load(f)
-
-# Load tokenizer - Corrected to use ByteLevelBPETokenizer with both files
-tokenizer = ByteLevelBPETokenizer("radiology_caption_model/vocab.json", "radiology_caption_model/merges.txt")
-
-# Instantiate the model with config parameters
-model = ImageCaptionModel(
-    vocab_size=config["vocab_size"],
-    embed_dim=config["embed_dim"]
-)
-
-# Load the model weights
-model.load_state_dict(torch.load("radiology_caption_model/pytorch_model.bin", map_location=torch.device('cpu')))
-model.eval() # Set model to evaluation mode
-
-# Define image transformations
-image_size = 128 # Must match training image size
-img_transforms = transforms.Compose([
-    transforms.Resize((image_size, image_size)),
-    transforms.ToTensor(),
-    transforms.Normalize(mean=[0.485,0.456,0.406],
-                         std=[0.229,0.224,0.225]),
-])
-
-# Function to generate caption
 def generate_caption(image):
-    # Preprocess image
-    img_tensor = img_transforms(image).unsqueeze(0) # Add batch dimension
-
-    # Generate caption
-    with torch.no_grad():
-        # Get image features
-        image_features = model.encoder(img_tensor)
-
-        # Start caption generation with BOS token
-        # (We assume BOS token ID is 2 from tokenizer training in cell 1)
-        # (Padding token ID is 0)
-        caption_tokens = [tokenizer.token_to_id("[BOS]")]
-        max_len = config["max_len"] if "max_len" in config else 40 # Use max_len from config, fallback to 40
-
-        for _ in range(max_len - 1): # -1 because BOS is already there
-            input_tokens = torch.tensor(caption_tokens).unsqueeze(0) # Add batch dimension
-            output = model.decoder(input_tokens, image_features)
-            last_token_logits = output[0, -1, :]
-            predicted_token_id = torch.argmax(last_token_logits).item()
-
-            caption_tokens.append(predicted_token_id)
-
-            # Stop if EOS token is generated
-            if predicted_token_id == tokenizer.token_to_id("[EOS]"):
-                break
-
-    # Decode the output tokens, excluding BOS and EOS (if present)
-    decoded_caption = tokenizer.decode(caption_tokens[1:-1] if caption_tokens[-1] == tokenizer.token_to_id("[EOS]") else caption_tokens[1:])
-    return decoded_caption
+    """
+    Generates a radiology caption for a given image
+    """
+    if isinstance(image, str):  # if image is a URL
+        image = Image.open(requests.get(image, stream=True).raw).convert("RGB")
+    elif isinstance(image, Image.Image):
+        image = image.convert("RGB")
+
+    inputs = processor(images=image, return_tensors="pt")
+    out = model.generate(**inputs)
+    caption = processor.decode(out[0], skip_special_tokens=True)
+    return caption
 
 # Create Gradio interface
-interface = gr.Interface(
+title = "Radiology Image Captioning"
+description = "Upload a radiology image (X-ray, CT, MRI) and get an automatic caption generated by the `hackergeek/radiology-image-captioning` model."
+
+iface = gr.Interface(
     fn=generate_caption,
-    inputs=gr.Image(type="pil"),
-    outputs="text",
-    title="Radiology Image Captioning",
-    description="Upload a radiology image (X-ray, CT, MRI) to get an AI-generated caption."
+    inputs=gr.Image(type="pil", label="Upload Radiology Image"),
+    outputs=gr.Textbox(label="Generated Caption"),
+    title=title,
+    description=description,
+    examples=[
+        ["https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/medical_xray.png"]
+    ]
 )
 
 if __name__ == "__main__":
-    interface.launch(server_name="0.0.0.0", server_port=7860)
+    iface.launch()