app.py

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=2,   # ✅ אפוקים מהירים: רק 2!
    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()