Update model/model.py

model/model.py

@@ -1,35 +1,36 @@
-from transformers import AutoModel, AutoModelForImageClassification
-from PIL import Image
-import numpy as np
 import torch
-import torch.nn as nn
-import torchvision.transforms as transforms
+from torchvision import transforms
+import torchvision.transforms.functional as F
+from transformers import AutoModelForImageClassification
 
 def predict(image_path):
     model = AutoModelForImageClassification.from_pretrained('sensei-ml/concrete_crack_images_classification')
     model.eval()
 
     with torch.no_grad():
-        # Define the transformations for the input image
-        transform = transforms.Compose([
-            transforms.Resize((227, 227)),
-            transforms.ToTensor(),
-            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
-        ])
-    
-        # Preprocess the image
-        image_tensor = torch.from_numpy(image_path).permute(2, 0, 1).float() 
-        input_tensor = transform(image_tensor).unsqueeze(0)  # Add batch dimension
-        
-        
+        # Convertir el array de NumPy a un tensor de PyTorch
+        image_tensor = torch.from_numpy(image_path).permute(2, 0, 1).float()  # Cambiar dimensiones de (H, W, C) a (C, H, W)
+
+        # Redimensionar la imagen usando funciones de transformación que soporten tensores
+        image_tensor = F.resize(image_tensor, [227, 227])
+
+        # Normalización
+        transform = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
+
+        # Aplicar la normalización
+        input_tensor = transform(image_tensor).unsqueeze(0)  # Añadir la dimensión del batch
+
+        # Hacer predicciones
         outputs = model(input_tensor)
         logits = outputs.logits
         probabilities = torch.nn.functional.softmax(logits, dim=1)[0]
         _, predicted_label = torch.max(logits, 1)
-        predicted_label = predicted_label.item()    
-        # Define the labels
+        predicted_label = predicted_label.item()
+
+        # Definir las etiquetas
         labels = ['Negative', 'Positive']
         probability_dict = {}
         for i, prob in enumerate(probabilities):
-          probability_dict[labels[i]] = prob.item()
+            probability_dict[labels[i]] = prob.item()
+
         return probability_dict