Create train.py

train.py

@@ -0,0 +1,73 @@
+# train.py
+from transformers import ViTImageProcessor, ViTForImageClassification, TrainingArguments, Trainer
+from transformers import DefaultDataCollator
+from datasets import load_dataset, Image
+import torch
+
+# 1. Charger le dataset et le mapper aux classes
+dataset = load_dataset("ashraq/fashion-product-images-small", name="styles", split="train")
+dataset = dataset.train_test_split(test_size=0.2)
+train_ds = dataset["train"]
+test_ds = dataset["test"]
+
+# 2. Créer la liste des labels (catégories uniques)
+labels = train_ds.unique("articleType")
+label2id, id2label = {}, {}
+for i, label in enumerate(labels):
+    label2id[label] = i
+    id2label[i] = label
+
+# 3. Charger le processeur et le modèle de base CORRECTS
+# On prend un modèle pré-entraîné sur ImageNet, pas sur des haricots !
+model_ckpt = "google/vit-base-patch16-224"
+processor = ViTImageProcessor.from_pretrained(model_ckpt)
+model = ViTForImageClassification.from_pretrained(
+    model_ckpt,
+    num_labels=len(labels),
+    id2label=id2label,
+    label2id=label2id,
+    ignore_mismatched_sizes=True # Important car le nombre de classes change
+)
+
+# 4. Fonction de preprocessing pour transformer les images
+def transform(example_batch):
+    inputs = processor([Image.open(img).convert("RGB") for img in example_batch["image_path"]], return_tensors="pt")
+    inputs["labels"] = [label2id[label] for label in example_batch["articleType"]]
+    return inputs
+
+# Appliquer le preprocessing
+train_ds = train_ds.cast_column("image_path", Image())
+test_ds = test_ds.cast_column("image_path", Image())
+
+train_ds.set_transform(transform)
+test_ds.set_transform(transform)
+
+# 5. Définir les arguments d'entraînement
+training_args = TrainingArguments(
+    output_dir="./vit-fashion-classifier",
+    per_device_train_batch_size=16,
+    evaluation_strategy="steps",
+    num_train_epochs=4,
+    fp16=True,
+    save_steps=100,
+    eval_steps=100,
+    logging_steps=10,
+    learning_rate=2e-4,
+    save_total_limit=2,
+    remove_unused_columns=False,
+    push_to_hub=True, # Pour pousser directement sur votre HF Space après l'entraînement
+    hub_model_id="MODLI/vit-fashion-classifier", # Remplacez par votre repo
+)
+
+# 6. Lancer l'entraînement
+trainer = Trainer(
+    model=model,
+    args=training_args,
+    data_collator=DefaultDataCollator(),
+    train_dataset=train_ds,
+    eval_dataset=test_ds,
+    tokenizer=processor,
+)
+
+trainer.train()
+trainer.push_to_hub()