Update app.py

app.py

@@ -5,56 +5,107 @@ 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
+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 SegmentationAlbumentationsTransform(ItemTransform):
+    split_idx = 0
+
+    def __init__(self, aug):
+        self.aug = aug
+
+    def encodes(self, x):
+        img,mask = x
+        aug = self.aug(image=np.array(img), mask=np.array(mask))
+        return PILImage.create(aug["image"]), PILMask.create(aug["mask"])
+
 class TargetMaskConvertTransform(ItemTransform):
     def __init__(self):
         pass
     def encodes(self, x):
-        img, mask = x
+        img,mask = x
+
+        #Convert to array
         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']
+        # 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
 repo_id = "LuisCe/Practica03"
 learner = from_pretrained_fastai(repo_id)
 
-# Define la función de predicción
-def predict_image(img):
-    img_fastai = Image.fromarray(img.astype('uint8'), 'RGB')
-    pred, _, _ = learner.predict(img_fastai)
-    return pred
+
+# 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)
+
+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)
+
+
+  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(predict_image, 
-              inputs="image", 
-              outputs="text", 
+gr.Interface(prediccion,
+              inputs="image",
+              outputs="image",
               title="Grape Segmentation",
               description="Segment grapes in the image.",
               theme="compact",
-              allow_flagging=False).launch()
+              allow_flagging=False).launch(debug=True)