guimCC/LORA_SemanticSegmentation

1import random2import gradio as gr3from transformers import SegformerForSemanticSegmentation, SegformerImageProcessor4from torchvision.transforms import ColorJitter, functional as F5from PIL import Image, ImageDraw, ImageFont6import numpy as np7import torch8from datasets import load_dataset9import evaluate10 11# Define the device12device = torch.device("cuda" if torch.cuda.is_available() else "cpu")13 14# Load the models15original_model_id = "guimCC/segformer-v0-gta"16lora_model_id = "guimCC/segformer-v0-gta-cityscapes"17 18original_model = SegformerForSemanticSegmentation.from_pretrained(original_model_id).to(device)19lora_model = SegformerForSemanticSegmentation.from_pretrained(lora_model_id).to(device)20 21# Load the dataset and select the first 10 images22dataset = load_dataset("Chris1/cityscapes", split="validation")23sampled_dataset = dataset.select(range(10))  # Select the first 10 examples24 25 26# Define your custom image processor27jitter = ColorJitter(brightness=0.25, contrast=0.25, saturation=0.25, hue=0.1)28 29# Initialize mIoU metric30metric = evaluate.load("mean_iou")31 32# Define id2label and processor if not already defined33id2label = {34    0: 'road', 1: 'sidewalk', 2: 'building', 3: 'wall', 4: 'fence', 5: 'pole',35    6: 'traffic light', 7: 'traffic sign', 8: 'vegetation', 9: 'terrain',36    10: 'sky', 11: 'person', 12: 'rider', 13: 'car', 14: 'truck', 15: 'bus',37    16: 'train', 17: 'motorcycle', 18: 'bicycle', 19: 'ignore'38}39processor = SegformerImageProcessor()40 41# Cityscapes color palette42palette = np.array([43    [128, 64, 128], [244, 35, 232], [70, 70, 70], [102, 102, 156], [190, 153, 153],44    [153, 153, 153], [250, 170, 30], [220, 220, 0], [107, 142, 35], [152, 251, 152],45    [70, 130, 180], [220, 20, 60], [255, 0, 0], [0, 0, 142], [0, 0, 70],46    [0, 60, 100], [0, 80, 100], [0, 0, 230], [119, 11, 32], [0, 0, 0]47])48 49def handle_grayscale_image(image):50    np_image = np.array(image)51    if np_image.ndim == 2:  # Grayscale image52        np_image = np.tile(np.expand_dims(np_image, -1), (1, 1, 3))53    return Image.fromarray(np_image)54 55def preprocess_image(image):56    image = handle_grayscale_image(image)57    image = jitter(image)  # Apply color jitter58    pixel_values = F.to_tensor(image).unsqueeze(0)  # Convert to tensor and add batch dimension59    return pixel_values.to(device)60 61def postprocess_predictions(logits):62    logits = logits.squeeze().detach().cpu().numpy()63    segmentation = np.argmax(logits, axis=0).astype(np.uint8)  # Convert to 8-bit integer64    return segmentation65 66def compute_miou(logits, labels):67    with torch.no_grad():68        logits_tensor = torch.from_numpy(logits)69        # Scale the logits to the size of the label70        logits_tensor = F.interpolate(71            logits_tensor,72            size=labels.shape[-2:],73            mode="bilinear",74            align_corners=False,75        ).argmax(dim=1)76 77        pred_labels = logits_tensor.detach().cpu().numpy()78        79        # Ensure the shapes of pred_labels and labels match80        if pred_labels.shape != labels.shape:81            labels = np.resize(labels, pred_labels.shape)82        83        pred_labels = [pred_labels]  # Wrap in a list84        labels = [labels]  # Wrap in a list85        86        metrics = metric.compute(87            predictions=pred_labels,88            references=labels,89            num_labels=len(id2label),90            ignore_index=19,91            reduce_labels=processor.do_reduce_labels,92        )93        94        mean_iou = metrics.get('mean_iou', 0.0)95        96        if np.isnan(mean_iou):97            mean_iou = 0.0  # Handle NaN values gracefully98        99        return mean_iou100    101def apply_color_palette(segmentation):102    colored_segmentation = palette[segmentation]103    return Image.fromarray(colored_segmentation.astype(np.uint8))104 105def create_legend():106    # Define font and its size107    try:108        font = ImageFont.truetype("arial.ttf", 15)109    except IOError:110        font = ImageFont.load_default()111 112    # Calculate legend dimensions113    num_classes = len(id2label)114    legend_height = 20 * ((num_classes + 1) // 2)  # Two items per row115    legend_width = 250116 117    # Create a blank image for the legend118    legend = Image.new("RGB", (legend_width, legend_height), (255, 255, 255))119    draw = ImageDraw.Draw(legend)120 121    # Draw each color and its label122    for i, (class_id, class_name) in enumerate(id2label.items()):123        color = tuple(palette[class_id])124        x = (i % 2) * 120125        y = (i // 2) * 20126        draw.rectangle([x, y, x + 20, y + 20], fill=color)127        draw.text((x + 30, y + 5), class_name, fill=(0, 0, 0), font=font)128 129    return legend130 131def inference(index, legend):132    """Run inference on the input image with both models."""133    image = sampled_dataset[index]['image']  # Fetch image from the sampled dataset134    pixel_values = preprocess_image(image)135 136    # Original model inference137    with torch.no_grad():138        original_outputs = original_model(pixel_values=pixel_values)139        original_segmentation = postprocess_predictions(original_outputs.logits)140 141    # LoRA model inference142    with torch.no_grad():143        lora_outputs = lora_model(pixel_values=pixel_values)144        lora_segmentation = postprocess_predictions(lora_outputs.logits)145 146    # Apply color palette147    original_segmentation_image = apply_color_palette(original_segmentation)148    lora_segmentation_image = apply_color_palette(lora_segmentation)149 150    # Return the original image, the segmentations, and mIoU151    return (152        image,153        original_segmentation_image,154        lora_segmentation_image,155    )156 157# Create a list of image options for the user to select from158image_options = [(f"Image {i}", i) for i in range(len(sampled_dataset))]159 160# Create the Gradio interface161iface = gr.Interface(162    fn=inference,163    inputs=[164        gr.Dropdown(label="Select Image", choices=image_options),165        gr.Image(type="pil", label="Legend", value=create_legend)166    ],167    outputs=[168        gr.Image(type="pil", label="Input Image"),169        gr.Image(type="pil", label="Original Model Prediction"),170        gr.Image(type="pil", label="LoRA Model Prediction"),171 172    ],173    live=True174)175 176# Launch the interface177iface.launch()178