Create train.py

train.py

@@ -0,0 +1,105 @@
+import torch
+import numpy as np
+from torchvision import transforms
+from datasets import load_dataset
+from transformers import (
+    SegformerForSemanticSegmentation,
+    SegformerFeatureExtractor,
+    Trainer,
+    TrainingArguments
+)
+import evaluate
+
+# ------------------------------
+# 1️⃣ Parameter
+# ------------------------------
+DATA_DIR = "path_to_dataset"  # Pfad zu deinen Bild- und Maskenordnern
+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(batch):
+    batch["pixel_values"] = [train_transforms(x) for x in batch["image"]]
+    # Masken als LongTensor für CrossEntropyLoss
+    batch["labels"] = [mask_transforms(x).long().squeeze(0) for x in batch["label"]]
+    return batch
+
+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!")