Spaces:
Running
Running
File size: 5,510 Bytes
4875d48 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 |
from dataclasses import dataclass
from typing import Optional, Tuple, Union, List
from PIL import Image
import PIL
import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import (
PreTrainedModel,
CLIPSegProcessor,
CLIPSegForImageSegmentation,
)
from transformers.modeling_outputs import ModelOutput
from .config import ClipSegMultiClassConfig
from sklearn.metrics import precision_score, recall_score, f1_score, accuracy_score
import numpy as np
from torch.utils.data import DataLoader
from collections import defaultdict
def flatten_outputs(preds, targets, num_classes):
"""Flatten predictions and targets to 1D arrays, filter ignored labels."""
preds = preds.cpu().numpy().reshape(-1)
targets = targets.cpu().numpy().reshape(-1)
mask = (targets >= 0) & (targets < num_classes)
return preds[mask], targets[mask]
def compute_metrics(all_preds, all_targets, num_classes, average="macro"):
y_pred = np.concatenate(all_preds)
y_true = np.concatenate(all_targets)
metrics = {
"accuracy": accuracy_score(y_true, y_pred),
"precision": precision_score(y_true, y_pred, average=average, zero_division=0),
"recall": recall_score(y_true, y_pred, average=average, zero_division=0),
"f1": f1_score(y_true, y_pred, average=average, zero_division=0),
}
return metrics
@dataclass
class ClipSegMultiClassOutput(ModelOutput):
loss: Optional[torch.FloatTensor] = None
logits: Optional[torch.FloatTensor] = None
predictions: Optional[torch.LongTensor] = None
class ClipSegMultiClassModel(PreTrainedModel):
config_class = ClipSegMultiClassConfig
base_model_prefix = "clipseg_multiclass"
def __init__(self, config: ClipSegMultiClassConfig):
super().__init__(config)
self.config = config
self.class_labels = config.class_labels
self.num_classes = config.num_classes
self.processor = CLIPSegProcessor.from_pretrained(config.model)
self.clipseg = CLIPSegForImageSegmentation.from_pretrained(config.model)
self.loss_fct = nn.CrossEntropyLoss()
def forward(
self,
pixel_values: Optional[torch.Tensor] = None,
input_ids: Optional[torch.Tensor] = None,
labels: Optional[torch.Tensor] = None,
**kwargs
) -> ClipSegMultiClassOutput:
if pixel_values is None or input_ids is None:
raise ValueError("Both `pixel_values` and `input_ids` must be provided.")
pixel_values = pixel_values.to(self.device)
input_ids = input_ids.to(self.device)
outputs = self.clipseg(pixel_values=pixel_values, input_ids=input_ids)
raw_logits = outputs.logits # shape: [B * C, H, W]
B = raw_logits.shape[0] // self.num_classes
C = self.num_classes
H, W = raw_logits.shape[-2:]
logits = raw_logits.view(B, C, H, W) # [B, C, H, W]
pred = torch.argmax(logits, dim=1) # [B, H, W]
loss = self.loss_fct(logits, labels.long()) if labels is not None else None
return ClipSegMultiClassOutput(
loss=loss,
logits=logits,
predictions=pred
)
@torch.no_grad()
def predict(self, images: Union[List, "PIL.Image.Image"]) -> torch.Tensor:
self.eval()
if isinstance(images, Image.Image):
images = [images]
inputs = self.processor(
images=[img for img in images for _ in self.class_labels],
text=self.class_labels * len(images),
return_tensors="pt",
padding=True,
truncation=True
).to(self.device)
output = self.forward(
pixel_values=inputs["pixel_values"],
input_ids=inputs["input_ids"]
)
return output.predictions
def evaluate(self, dataloader: torch.utils.data.DataLoader) -> dict:
from sklearn.metrics import accuracy_score, precision_score, recall_score, f1_score
import numpy as np
self.eval()
all_preds = []
all_targets = []
with torch.no_grad():
for batch in dataloader:
pixel_values = batch["pixel_values"].to(self.device) # [B * C, 3, H, W]
input_ids = batch["input_ids"].to(self.device) # [B * C, T]
labels = batch["labels"].to(self.device) # [B, H, W]
outputs = self.forward(pixel_values=pixel_values, input_ids=input_ids)
preds = outputs.predictions # [B, H, W]
for pred, label in zip(preds, labels):
pred = pred.cpu().flatten()
label = label.cpu().flatten()
mask = label != 0
pred = pred[mask]
label = label[mask]
all_preds.append(pred)
all_targets.append(label)
y_pred = torch.cat(all_preds).numpy()
y_true = torch.cat(all_targets).numpy()
return {
"accuracy": accuracy_score(y_true, y_pred),
"precision": precision_score(y_true, y_pred, average="macro", zero_division=0),
"recall": recall_score(y_true, y_pred, average="macro", zero_division=0),
"f1": f1_score(y_true, y_pred, average="macro", zero_division=0),
}
|