Update tasks/image.py

tasks/image.py

@@ -3,6 +3,7 @@ import torch.nn as nn
 import torch.optim as optim
 from torchvision import transforms
 from torch.utils.data import DataLoader, Dataset
+from huggingface_hub import hf_hub_download
 
 from fastapi import APIRouter
 from datetime import datetime
@@ -180,19 +181,30 @@ async def evaluate_image(request: ImageEvaluationRequest):
 
 
     # Training loop 
-    num_epochs = 10 
-    for epoch in range(num_epochs):
-        for images, labels in train_loader :
-            images, labels = images.to(device), labels.to(device)
-            # Zero the parameter gradients
-            optimizer.zero_grad()
+    # num_epochs = 10 
+    # for epoch in range(num_epochs):
+    #     for images, labels in train_loader :
+    #         images, labels = images.to(device), labels.to(device)
+    #         # Zero the parameter gradients
+    #         optimizer.zero_grad()
 
-            # Forward + backward + optimize
-            outputs = model(images)
-            loss = criterion(outputs, labels)
-            loss.backward()
-            optimizer.step()
-        print(f'Epoch [{epoch + 1}/10], Loss: {loss.item():.4f}')
+    #         # Forward + backward + optimize
+    #         outputs = model(images)
+    #         loss = criterion(outputs, labels)
+    #         loss.backward()
+    #         optimizer.step()
+    #     print(f'Epoch [{epoch + 1}/10], Loss: {loss.item():.4f}')
+
+    # Charging pre-trained model
+    repo_id = "AlexandreL2024/CNN-Image-Classification"
+    filename = "model_CNN_2Layers.pth"
+    
+    # Upload file .pth from Hugging Face
+    model_path = hf_hub_download(repo_id=repo_id, filename=filename)
+    
+    # Charger le modèle avec torch.load()
+    model = ImageClassifier()
+    model = model.load_state_dict(torch.load(model_path))
 
     # Evaluation loop
     model.eval()  # Set the model to evaluation mode