Create app.py

app.py

@@ -0,0 +1,63 @@
+import gradio as gr
+from datasets import load_dataset
+from transformers import AutoImageProcessor, AutoModelForImageClassification, Trainer, TrainingArguments
+import torch
+import os
+
+# טוענים דאטאסט מהתיקיות
+dataset = load_dataset("imagefolder", data_dir=".", split={"train": "train[:80%]", "test": "train[80%:]"})
+
+# בוחרים מודל בסיסי
+checkpoint = "google/vit-tiny-patch16-224"
+processor = AutoImageProcessor.from_pretrained(checkpoint)
+model = AutoModelForImageClassification.from_pretrained(
+    checkpoint,
+    num_labels=3,
+    id2label={0: "rock", 1: "paper", 2: "scissors"},
+    label2id={"rock": 0, "paper": 1, "scissors": 2}
+)
+
+# פונקציה לעיבוד התמונות
+def preprocess(examples):
+    images = [x.convert("RGB") for x in examples["image"]]
+    inputs = processor(images=images, return_tensors="pt")
+    inputs["labels"] = examples["label"]
+    return inputs
+
+dataset = dataset.map(preprocess, batched=True)
+
+# הגדרות אימון
+training_args = TrainingArguments(
+    output_dir="./results",
+    evaluation_strategy="epoch",
+    save_strategy="epoch",
+    per_device_train_batch_size=4,
+    per_device_eval_batch_size=4,
+    num_train_epochs=5,
+    load_best_model_at_end=True,
+    logging_dir='./logs',
+    logging_steps=5,
+)
+
+trainer = Trainer(
+    model=model,
+    args=training_args,
+    train_dataset=dataset["train"],
+    eval_dataset=dataset["test"],
+)
+
+# אימון
+trainer.train()
+
+# פונקציה להרצת חיזוי על תמונה חדשה
+def predict(image):
+    inputs = processor(images=image, return_tensors="pt")
+    outputs = model(**inputs)
+    logits = outputs.logits
+    predicted_class_idx = logits.argmax(-1).item()
+    label = model.config.id2label[predicted_class_idx]
+    return label
+
+# בניית אפליקציה
+demo = gr.Interface(fn=predict, inputs="image", outputs="text")
+demo.launch()