import torch
import torch.nn as nn
import torch.optim as optim
from definition import FlowersImagesDetectionModel
from torch.utils.data import DataLoader
from datasets import load_dataset
from torchvision.transforms import ToTensor, Resize
from torch.utils.data.dataset import TensorDataset

flowerTypesNumber = 102

model = FlowersImagesDetectionModel(flowerTypesNumber)

# Funzioni di ottimizzazione e di perdita
optimizer = optim.Adam(model.parameters(), lr=0.001)
criterion = nn.CrossEntropyLoss()

# Caricamento del dataset
originalDataset = load_dataset("nelorth/oxford-flowers", split="train")

tensorImages = []
tensorLabels = []

# Trasforma le immagini in tensori PyTorch e ridimensionale
for imageData, label in zip(originalDataset['image'], originalDataset['label']):
    tensorImage = ToTensor()(Resize((224, 224))(imageData))  # Ridimensiona le immagini
    tensorImages.append(tensorImage)
    tensorLabels.append(label)

# Trasforma le liste di tensori in un singolo tensore
imagesTensor = torch.stack(tensorImages)
labelsTensor = torch.tensor(tensorLabels)

# Crea un dataset
dataset = TensorDataset(imagesTensor, labelsTensor)

# Crea un DataLoader
dataLoader = DataLoader(dataset, batch_size=64, shuffle=True)

# Addestramento
model.train()
for epoch in range(2):
    running_loss = 0.0

    for i, (inputs, labels) in enumerate(dataLoader, 0):
        optimizer.zero_grad()
        outputs = model(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()
        running_loss += loss.item()

        if i % 100 == 99:
            print('[%d, %5d] loss: %.3f' %
                  (epoch + 1, i + 1, running_loss / 100))
            running_loss = 0.0

torch.save(model.state_dict(), 'pytorch_model.bin')