Update app.py

app.py

@@ -1,76 +1,59 @@
 import torch
-from torchvision import transforms
-from PIL import Image, ImageStat
-import gradio as gr
+import torch.nn as nn
+from torchvision import models, transforms
+from PIL import Image
 import numpy as np
+import gradio as gr
 
-# Carica il tuo modello pre-addestrato (aggiorna il percorso del modello come necessario)
-model = torch.load('model_v11.pt', map_location=torch.device('cpu'))
-model.eval()
+# Definisci il modello e il numero di classi
+class PretrainedModel(nn.Module):
+    def __init__(self, num_classes=19):
+        super(PretrainedModel, self).__init__()
+        
+        weights = models.ResNet50_Weights.IMAGENET1K_V2
+        self.model = models.resnet50(weights=weights)
+        
+        # Congela i layer iniziali
+        for param in self.model.parameters():
+            param.requires_grad = False
+            
+        # Sostituisci l'ultimo layer
+        num_features = self.model.fc.in_features
+        self.model.fc = nn.Linear(num_features, num_classes)
+    
+    def forward(self, x):
+        return self.model(x)
 
-# Funzione per loggare i colori dell'immagine
-def log_image_color_stats(image):
-    stat = ImageStat.Stat(image)
-    mean_rgb = stat.mean
-    print(f"Color Stats - R: {mean_rgb[0] / 255.0}, G: {mean_rgb[1] / 255.0}, B: {mean_rgb[2] / 255.0}")
+# Crea un'istanza del modello
+model = PretrainedModel(num_classes=19)
 
-# Funzione per loggare i dati del tensore
-def log_tensor_data(tensor):
-    print(f"Tensor shape: {tensor.shape}")
-    print(f"Tensor data (first 10 values): {tensor.flatten()[:10]}")
+# Carica i pesi
+model.load_state_dict(torch.load('model_v11.pt', map_location=torch.device('cpu')))
+model.eval()  # Imposta il modello in modalità valutazione
 
-# Preprocessing per l'immagine con normalizzazione
+# Trasformazioni
 preprocess = transforms.Compose([
     transforms.Resize((224, 224)),
     transforms.ToTensor(),
     transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
 ])
 
-# Preprocessing senza normalizzazione (separato)
-preprocess_without_normalize = transforms.Compose([
-    transforms.Resize((224, 224)),
-    transforms.ToTensor()
-])
-
-# Funzione di previsione
-def classify_image(image):
-    # Logga i colori dell'immagine originale
-    log_image_color_stats(image)
-    
-    # Ridimensiona l'immagine per coerenza
-    img_resized = image.resize((224, 224))
-    print(f"Resized Bitmap dimensions: {img_resized.size}")
-
-    # Preprocessa e ottieni il tensore non normalizzato
-    img_tensor_no_normalize = preprocess_without_normalize(image).unsqueeze(0)
-    print("Tensor data before normalization:")
-    log_tensor_data(img_tensor_no_normalize)
-    
-    # Preprocessa e normalizza
-    img_tensor = preprocess(image).unsqueeze(0)
-    log_tensor_data(img_tensor)
-
-    # Esegui la previsione
+def classify_image(img):
+    # Preprocessa l'immagine
+    img_tensor = preprocess(img).unsqueeze(0)  # Aggiunge una dimensione per il batch
     with torch.no_grad():
         output = model(img_tensor)
-        
-        # Logga i punteggi raw
-        print(f"Raw output scores: {output}")
-
-        # Calcola le probabilità
         probabilities = torch.softmax(output, dim=1)
-
-        # Estrai le probabilità e l'indice della classe prevista
-        predicted_class_index = probabilities.argmax()
-        predicted_probability = probabilities[0][predicted_class_index]
+        predicted_class_index = probabilities.argmax().item()
+        predicted_probability = probabilities[0][predicted_class_index].item()
         
-        # Logga le probabilità di ciascuna classe
-        print("Probabilities for each class:")
-        for index, prob in enumerate(probabilities[0]):
-            print(f"Class {index}: {prob.item():.4f}")
+    return f"Class {predicted_class_index}, Confidence: {predicted_probability:.4f}"
 
-    # Restituisci l'indice della classe e la probabilità come output
-    return f"Predicted class index: {predicted_class_index.item()}, Probability: {predicted_probability.item():.4f}"
+# Configura Gradio
+iface = gr.Interface(
+    fn=classify_image,
+    inputs=gr.inputs.Image(type="pil"),
+    outputs="text"
+)
 
-# Interfaccia Gradio
-gr.Interface(fn=classify_image, inputs="image", outputs="text").launch()
+iface.launch()