Update app.py

app.py

@@ -1,107 +1,125 @@
+import gradio as gr
 import torch
-import numpy as np
+from PIL import Image
 from torchvision import transforms
-from datasets import load_dataset
-from transformers import (
-    SegformerForSemanticSegmentation,
-    SegformerFeatureExtractor,
-    Trainer,
-    TrainingArguments
-)
+import numpy as np
+from matplotlib import pyplot as plt
+from torch.utils.data import Dataset, DataLoader
 import evaluate
-
-# ------------------------------
-# 1️⃣ Parameter
-# ------------------------------
-DATA_DIR = "path_to_dataset"  
-NUM_CLASSES = 3               # z.B. 3 Klassen: Hintergrund, Schaden, Rand
-IMAGE_SIZE = 256              # Bildgröße für Training
-
-# ------------------------------
-# 2️⃣ Dataset laden
-# ------------------------------
-# Annahme: Dataset im ImageFolder Format mit Unterordnern 'train' und 'validation'
-dataset = load_dataset("imagefolder", data_dir=DATA_DIR)
-
-# Transformationen für Bilder
-train_transforms = transforms.Compose([
-    transforms.Resize((IMAGE_SIZE, IMAGE_SIZE)),
-    transforms.ToTensor(),
-])
-
-mask_transforms = transforms.Compose([
-    transforms.Resize((IMAGE_SIZE, IMAGE_SIZE)),
-    transforms.PILToTensor(),  # Masken als Tensor
-])
-
-# Preprocessing-Funktion
-def preprocess(example):
-    example["pixel_values"] = train_transforms(example["image"])
-    # Masken als LongTensor für CrossEntropyLoss
-    example["labeks"] = mask_transforms(example["label"]).long().squeeze(0) 
-    return example
-
-dataset = dataset.map(preprocess)
-
-# ------------------------------
-# 3️⃣ Feature Extractor & Modell
-# ------------------------------
-feature_extractor = SegformerFeatureExtractor.from_pretrained("nvidia/mit-b1")
-
-model = SegformerForSemanticSegmentation.from_pretrained(
-    "nvidia/mit-b1",
-    num_labels=NUM_CLASSES,
-)
-
-# ------------------------------
-# 4️⃣ Metrics
-# ------------------------------
-metric = evaluate.load("mean_iou")
-
-def compute_metrics(p):
-    preds = np.argmax(p.predictions, axis=1)
-    return metric.compute(predictions=preds, references=p.label_ids, num_labels=NUM_CLASSES)
-
-# ------------------------------
-# 5️⃣ TrainingArguments
-# ------------------------------
-training_args = TrainingArguments(
-    output_dir="./results",
-    per_device_train_batch_size=4,
-    per_device_eval_batch_size=4,
-    num_train_epochs=10,
-    learning_rate=5e-5,
-    evaluation_strategy="steps",
-    save_strategy="steps",
-    save_steps=200,
-    eval_steps=200,
-    logging_steps=50,
-    fp16=True,  # Mixed Precision, falls GPU verfügbar
-    remove_unused_columns=False,  # wichtig für Segmentation
-)
-
-# ------------------------------
-# 6️⃣ Trainer
-# ------------------------------
-trainer = Trainer(
-    model=model,
-    args=training_args,
-    train_dataset=dataset["train"],
-    eval_dataset=dataset["validation"],
-    compute_metrics=compute_metrics,
-)
-
-# ------------------------------
-# 7️⃣ Training starten
-# ------------------------------
-trainer.train()
-
-# ------------------------------
-# 8️⃣ Modell speichern
-# ------------------------------
-trainer.save_model("my_segformer_model")
-feature_extractor.save_pretrained("my_segformer_model")
-
-print("✅ Training abgeschlossen und Modell gespeichert!")
-
-n
+from torch import nn
+from transformers import SegformerForSemanticSegmentation
+import sys
+import io 
+
+
+###################
+# Setup label names
+target_list = ['Crack', 'ACrack', 'Wetspot', 'Efflorescence', 'Rust', 'Rockpocket', 'Hollowareas', 'Cavity',
+               'Spalling', 'Graffiti', 'Weathering', 'Restformwork', 'ExposedRebars', 
+               'Bearing', 'EJoint', 'Drainage', 'PEquipment', 'JTape', 'WConccor']
+classes, nclasses = target_list, len(target_list)
+label2id = dict(zip(classes, range(nclasses)))
+id2label = dict(zip(range(nclasses), classes))
+
+############
+# Load model
+device = torch.device('cpu')
+segformer = SegformerForSemanticSegmentation.from_pretrained("nvidia/mit-b1",
+            id2label=id2label,
+            label2id=label2id)
+
+# SegModel
+class SegModel(nn.Module):
+    def __init__(self, segformer):
+        super(SegModel, self).__init__()
+        self.segformer = segformer
+        self.upsample = nn.Upsample(scale_factor=4, mode='nearest')
+
+    def forward(self, x):
+        return self.upsample(self.segformer(x).logits)
+ 
+model = SegModel(segformer)  
+path = "runs/2023-08-31_rich-paper-12/best_model_cpu.pth"
+print(f"Load Segformer weights from {path}")
+#model = model.load_state_dict(torch.load(path, map_location=device))
+model = torch.load(path)
+model.eval()
+
+##################
+# Image preprocess
+##################
+
+to_tensor = transforms.ToTensor()
+resize = transforms.Resize((512, 512))
+normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
+                                     std=[0.229, 0.224, 0.225])
+
+def process_pil(img):
+    img = to_tensor(img)
+    img = resize(img)
+    img = normalize(img)
+    return img
+
+###########
+# Inference
+
+def inference(img, name):
+    img = process_pil(img)
+    mask = model(img.unsqueeze(0)) # we need a batch, hence we introduce an extra dimenation at position 0 (unsqueeze)
+    mask = mask[0]
+
+    # Get probability values (logits to probs)
+    mask_probs = torch.sigmoid(mask)
+    mask_probs = mask_probs.detach().numpy()
+    mask_probs.shape
+
+    # Make binary mask
+    THRESHOLD = 0.5
+    mask_preds = mask_probs > THRESHOLD
+
+    # All combined
+    mask_all =  mask_preds.sum(axis=0)
+    mask_all = np.expand_dims(mask_all, axis=0)
+    mask_all.shape
+
+    # Concat all combined with normal preds
+    mask_preds = np.concatenate((mask_all, mask_preds),axis=0)
+    labs = ["ALL"] + target_list
+
+    fig, axes = plt.subplots(5, 4, figsize = (10,10))
+            
+    for i, ax in enumerate(axes.flat):
+        label = labs[i]
+        ax.imshow(mask_preds[i])
+        ax.set_title(label)
+          
+    plt.tight_layout()        
+
+
+    # plt to PIL
+    img_buf = io.BytesIO()
+    fig.savefig(img_buf, format='png')
+    im = Image.open(img_buf)
+    return im
+
+
+
+title = "Masterarbeit"
+description = """
+
+
+"""
+
+article = "<p style='text-align: center'><a href='https://github.com/phiyodr/dacl10k-toolkit' target='_blank'>Github Repo</a></p>"
+examples=[['assets/dacl10k_v2_validation_0037.jpg', 'dacl10k_v2_validation_0037.jpg'],['assets/dacl10k_v2_validation_0068.jpg','dacl10k_v2_validation_0068.jpg'], ['assets/dacl10k_v2_validation_0053.jpg', 'dacl10k_v2_validation_0053.jpg']]
+
+demo = gr.Interface(
+    fn=inference, 
+    inputs=gr.inputs.Image(type="pil"), 
+    outputs=gr.outputs.Image(type="pil"),
+    title=title,
+    description=description,
+    article=article,
+    examples=examples)
+
+demo.launch()