app.py

import gradio as gr
from datasets import load_dataset
from PIL import Image, ImageDraw
import numpy as np
from rdp import rdp

# Load the dataset
dataset = load_dataset("dwb2023/brain-tumor-image-dataset-semantic-segmentation", split="test")

def simplify_segmentation(segmentation, max_points=20):
    if not segmentation or len(segmentation) == 0:
        return []
    epsilon = 1.0
    simplified = rdp(np.array(segmentation), epsilon=epsilon)
    while len(simplified) > max_points:
        epsilon *= 1.5
        simplified = rdp(np.array(segmentation), epsilon=epsilon)
    return simplified.tolist()

def draw_annotations(index):
    try:
        # Fetch the image and annotations from the dataset
        record = dataset[index]
        
        # Convert image to PIL Image if it's a numpy array
        if isinstance(record['image'], np.ndarray):
            img = Image.fromarray(record['image'])
        else:
            img = record['image']
        
        img = img.convert("RGB")  # Ensure the image is in RGB mode
        
        draw = ImageDraw.Draw(img)

        # Draw bounding box
        bbox = record["bbox"]
        draw.rectangle([bbox[0], bbox[1], bbox[0] + bbox[2], bbox[1] + bbox[3]], outline="red", width=2)

        # Draw original segmentation mask
        segmentation = record["segmentation"]
        original_points = 0
        simplified_points = 0
        for seg in segmentation:
            if seg:  # Check if the segmentation is not empty
                draw.polygon(seg, outline="blue", width=2)
                original_points += len(seg)
                
                # Simplify and draw simplified segmentation
                simplified_seg = simplify_segmentation(seg)
                draw.polygon(simplified_seg, outline="green", width=2)
                simplified_points += len(simplified_seg)

        # Prepare additional information
        category_id = record["category_id"]
        area = record["area"]
        file_name = record["file_name"]

        info = f"File Name: {file_name}\n"
        info += f"Image ID: {record['id']}\n"
        info += f"Category ID: {category_id}\n"
        info += f"Bounding Box: [{bbox[0]:.2f}, {bbox[1]:.2f}, {bbox[2]:.2f}, {bbox[3]:.2f}]\n"
        info += f"Original Segmentation Points: {original_points}\n"
        info += f"Simplified Segmentation Points: {simplified_points}\n"
        info += f"Area: {area:.2f}"

        return img, info
    except Exception as e:
        print(f"Error processing image at index {index}: {e}")
        return Image.new('RGB', (300, 300), color='gray'), f"Error loading image information: {str(e)}"

# Create Gradio interface
with gr.Blocks() as demo:
    gr.Markdown("# Brain Tumor Image Dataset Viewer")
    gr.Markdown("## Refer to the [dwb2023/brain-tumor-image-dataset-semantic-segmentation](https://huggingface.co/datasets/dwb2023/brain-tumor-image-dataset-semantic-segmentation/viewer/default/test) dataset for more information")
    gr.Markdown("### Red: Bounding Box, Blue: Original Segmentation, Green: Simplified Segmentation (max 20 points)")
    
    with gr.Row():
        with gr.Column(scale=1):
            image_output = gr.Image(label="Annotated Image")
        with gr.Column(scale=1):
            image_index = gr.Slider(minimum=0, maximum=len(dataset)-1, step=1, value=0, label="Image ID Slider")
            info_output = gr.Textbox(label="Image Information", lines=10)
    
    # Update image and info when slider changes
    image_index.change(draw_annotations, inputs=image_index, outputs=[image_output, info_output])

    # Display initial image and info
    demo.load(draw_annotations, inputs=image_index, outputs=[image_output, info_output])

demo.launch(debug=True)