Update app.py

app.py

@@ -1,24 +1,90 @@
 import gradio as gr
+from fastai.vision.all import *
 from fastai.learner import load_learner
 from PIL import Image
 
-# Load the learner
-learn_inf = load_learner("LuisCe/Practica03.pkl")
-
-# Define the prediction function
-def predict_image(img):
-    # Convert the PIL image to a format that fastai expects
-    img_fastai = Image.fromarray(img.astype('uint8'), 'RGB')
-    # Make prediction
-    pred, _, _ = learn_inf.predict(img_fastai)
-    # Return prediction
-    return pred
-
-# Create Gradio interface
-gr.Interface(predict_image, 
-              inputs="image", 
-              outputs="text", 
+# Define la función get_y_fn si se utiliza en tu modelo
+def get_y_fn(x):
+    return Path(str(x).replace("Images","Labels").replace("color","gt").replace(".jpg",".png"))
+
+# Define la clase TargetMaskConvertTransform
+class TargetMaskConvertTransform(ItemTransform):
+    def __init__(self):
+        pass
+
+    def encodes(self, x):
+        img, mask = x
+        mask = np.array(mask)
+        mask[mask == 255] = 1
+        mask[mask == 150] = 2
+        mask[mask == 74] = 3
+        mask[mask == 76] = 3
+        mask[mask == 29] = 4
+        mask[mask == 25] = 4
+        mask = PILMask.create(mask)
+        return img, mask
+
+# Define la clase SegmentationAlbumentationsTransform
+class SegmentationAlbumentationsTransform(ItemTransform):
+    def __init__(self):
+        pass
+
+    def encodes(self, x):
+        img, mask = x
+        # Aquí deberías definir tu transformación de Albumentations
+        # Por ejemplo, podrías tener algo como:
+        # transformed = my_albumentations_function(image=img, mask=mask)
+        # return transformed['image'], transformed['mask']
+        return img, mask
+
+# Carga el modelo después de definir la clase
+repo_id = "LuisCe/Practica03"
+learner = from_pretrained_fastai(repo_id)
+
+
+# Carga el modelo previamente entrenado
+model = learner.model
+model = model.cpu()
+model.eval()
+
+import torchvision.transforms as transforms
+def transform_image(image):
+    my_transforms = transforms.Compose([transforms.ToTensor(),
+                                        transforms.Normalize(
+                                            [0.485, 0.456, 0.406],
+                                            [0.229, 0.224, 0.225])])
+    image_aux = image
+    return my_transforms(image_aux).unsqueeze(0).to(device)
+
+def prediccion(img):
+  image = transforms.Resize((480,640))(img)
+  tensor = transform_image(image=image)
+
+
+  model.to(device)
+  with torch.no_grad():
+      outputs = model(tensor)
+
+  outputs = torch.argmax(outputs,1)
+
+
+
+  mask = np.array(outputs.cpu())
+  mask[mask==1]=255
+  mask[mask==2]=150
+  mask[mask==3]=74
+  mask[mask==4]=29
+
+  mask=np.reshape(mask,(480,640))
+
+  return(mask)
+
+# Crea la interfaz Gradio
+gr.Interface(prediccion,
+              inputs="image",
+              outputs="image",
               title="Grape Segmentation",
               description="Segment grapes in the image.",
               theme="compact",
               allow_flagging=False).launch()
+