Update app.py

app.py

@@ -1,4 +1,3 @@
-from huggingface_hub import from_pretrained_fastai
 import gradio as gr
 from fastai.basics import *
 from fastai.vision import models
@@ -6,65 +5,12 @@ 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 transformss
-import torch
-from albumentations import *
-
-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
-
-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"])
+import torchvision.transforms as 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
-        
-        #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
-        
-        # Back to PILMask
-        mask = PILMask.create(mask)
-        return img, mask
-
-repo_id = "paascorb/practica3_Segmentation"
-
-learner = from_pretrained_fastai(repo_id)
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu") 
+model = torch.jit.load("unet.pth")
+model = model.cpu()
 
 def transform_image(image):
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
@@ -79,8 +25,7 @@ def predict(img):
     img = PIL.Image.fromarray(img, "RGB")
     image = transformss.Resize((480,640))(img)
     tensor = transform_image(image=image) 
-
-    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    
     learner.to(device)
     with torch.no_grad():
         outputs = learner(tensor)