Update app.py

app.py

@@ -3,38 +3,54 @@ from fastai.vision.all import *
 from fastai.learner import load_learner
 from PIL import Image
 
-# Define la función get_y_fn si se utiliza en tu modelo
-def get_y_fn(x):
+from albumentations import (
+    Compose,
+    OneOf,
+    ElasticTransform,
+    GridDistortion,
+    OpticalDistortion,
+    HorizontalFlip,
+    Rotate,
+    Transpose,
+    CLAHE,
+    ShiftScaleRotate
+)
+
+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
+class SegmentationAlbumentationsTransform(ItemTransform):
+    split_idx = 0
+
+    def __init__(self, aug):
+        self.aug = aug
 
     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
+        img,mask = x
+        aug = self.aug(image=np.array(img), mask=np.array(mask))
+        return PILImage.create(aug["image"]), PILMask.create(aug["mask"])
 
-# Define la clase SegmentationAlbumentationsTransform
-class SegmentationAlbumentationsTransform(ItemTransform):
+class TargetMaskConvertTransform(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']
+        img,mask = x
+
+        #Convert to array
+        mask = np.array(mask)
+
+        # mask[mask!=255]=0
+        # Change 255 for 1
+        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[mask==255]=1
+
+        # Back to PILMask
+        mask = PILMask.create(mask)
         return img, mask
 
 # Carga el modelo después de definir la clase
@@ -56,7 +72,11 @@ def transform_image(image):
     image_aux = image
     return my_transforms(image_aux).unsqueeze(0).to(device)
 
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+
 def prediccion(img):
+  img = Image.fromarray(img)
   image = transforms.Resize((480,640))(img)
   tensor = transform_image(image=image)
 
@@ -86,5 +106,4 @@ gr.Interface(prediccion,
               title="Grape Segmentation",
               description="Segment grapes in the image.",
               theme="compact",
-              allow_flagging=False).launch()
-
+              allow_flagging=False).launch(debug=True)