Update app.py

app.py

@@ -1,74 +1,80 @@
 from huggingface_hub import from_pretrained_fastai
 import gradio as gr
+from fastai.basics import *
+from fastai.vision import models
 from fastai.vision.all import *
+from fastai.metrics import *
+from fastai.data.all import *
+from fastai.callback import *
+
+
+from pathlib import Path
+import random
 import PIL
 import torchvision.transforms as transforms
-# from albumentations import (
-#     Compose,
-#     OneOf,
-#     ElasticTransform,
-#     GridDistortion, 
-#     OpticalDistortion,
-#     HorizontalFlip,
-#     Rotate,
-#     Transpose,
-#     CLAHE,
-#     ShiftScaleRotate
-# )
+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"))
+def get_y_fn (x):
+    return Path(str(x).replace("Images","Labels").replace("color","gt").replace(".jpg",".png"))
 
-# def ParentSplitter(x):
-#     return Path(x).parent.name==test_name
+def ParentSplitter(x):
+    return Path(x).parent.name==test_name
 
-# class SegmentationAlbumentationsTransform(ItemTransform):
-#     split_idx = 0
+class SegmentationAlbumentationsTransform(ItemTransform):
+    split_idx = 0
     
-#     def __init__(self, aug): 
-#         self.aug = aug
+    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"])
+    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"])
 
-# transforms=Compose([HorizontalFlip(p=0.5),
-#                     Rotate(p=0.40,limit=10),GridDistortion()
-#                     ],p=1)
-# transformPipeline=SegmentationAlbumentationsTransform(transforms)
+transforms=Compose([HorizontalFlip(p=0.5),
+                    Rotate(p=0.40,limit=10),GridDistortion()
+                    ],p=1)
+transformPipeline=SegmentationAlbumentationsTransform(transforms)
 
-# class TargetMaskConvertTransform(ItemTransform):
-#     def __init__(self): 
-#         pass
-#     def encodes(self, x):
-#         img,mask = x
+class TargetMaskConvertTransform(ItemTransform):
+    def __init__(self): 
+        pass
+    def encodes(self, x):
+        img,mask = x
         
-#         #Convert to array
-#         mask = np.array(mask)
+        #Convert to array
+        mask = np.array(mask)
         
-#         # Aquí definimos cada clase en la máscara
-#         # uva:
-#         mask[mask==255]=1
-#         # hojas:
-#         mask[mask==150]=2
-#         # conductores:
-#         mask[mask==76]=3
-#         mask[mask==74]=3
-#         # madera:
-#         mask[mask==29]=4
-#         mask[mask==25]=4
+        # Aquí definimos cada clase en la máscara
+        # uva:
+        mask[mask==255]=1
+        # hojas:
+        mask[mask==150]=2
+        # conductores:
+        mask[mask==76]=3
+        mask[mask==74]=3
+        # madera:
+        mask[mask==29]=4
+        mask[mask==25]=4
         
-#         # Back to PILMask
-#         mask = PILMask.create(mask)
-#         return img, mask
-
-# repo_id = "paascorb/practica3_Segmentation"
+        # Back to PILMask
+        mask = PILMask.create(mask)
+        return img, mask
 
-# learner = from_pretrained_fastai(repo_id)
+repo_id = "paascorb/practica3_Segmentation"
 
-learn = unet_learner(None,resnet18,cbs=None,metrics=[DiceMulti()]).to_fp16()
-learn.load('model')
+learner = from_pretrained_fastai(repo_id)
 
 def transform_image(image):
     my_transforms = transforms.Compose([transforms.ToTensor(),