Upload 2 files

aimodel.py

@@ -6,6 +6,8 @@ import uuid
 import re
 from fastai.vision.all import *
 from pathlib import Path
+import torch
+import torch.nn as nn
 import random
 import shutil
 
@@ -94,16 +96,23 @@ def train_model():
         batch_tfms=Normalize.from_stats(*imagenet_stats)
     )
     dls = dblock.dataloaders(dataset_path)
-    learn = vision_learner(dls, resnet18, metrics=accuracy) # Criar um novo Learner
+    num_classes = len(dls.vocab)
+    learn = vision_learner(dls, resnet18, metrics=accuracy, pretrained=True) # Começar com pesos pré-treinados
 
     if os.path.exists(utils.MODEL_PATH):
         try:
-            learn.load(utils.MODEL_PATH.replace(".pkl", "")) # Carregar apenas os pesos (sem o otimizador e outros estados)
-            logging.info(f"Existing model weights loaded from {utils.MODEL_PATH}")
+            # Carregar apenas os pesos do backbone (todas as camadas exceto a última linear)
+            state = torch.load(utils.MODEL_PATH)
+            # Remover a chave da camada de classificação (geralmente '1.fc' ou similar)
+            new_state = {k: v for k, v in state['model'].items() if not k.endswith('.weight') and not k.endswith('.bias')}
+            learn.model.load_state_dict(new_state, strict=False)
+            logging.info("Backbone weights loaded.")
         except Exception as e:
-            logging.error(f"Error loading existing model weights: {e}. Training from scratch.")
-    else:
-        logging.info("No existing model found. Training from scratch.")
+            logging.error(f"Error loading backbone weights: {e}. Training from scratch.")
+
+    # Substituir a camada de classificação
+    num_ftrs = learn.model.fc.in_features
+    learn.model.fc = nn.Linear(num_ftrs, num_classes)
 
     learn.fine_tune(5)
     learn.export(utils.MODEL_PATH)

requirements.txt

@@ -5,4 +5,5 @@ requests
 numpy
 opencv-python-headless
 python-multipart
-fastai
+fastai
+torch