Update app.py

app.py

@@ -4,23 +4,12 @@ from fastai.vision.all import *
 import PIL
 import torchvision.transforms as transforms
 
-class TargetMaskConvertTransform(ItemTransform):
-    def __init__(self): 
-        pass
-    def encodes(self, x):
-        img,mask = x
-        
-        #Convert to array
-        mask = np.array(mask)
-        
-        mask[mask!=255]=0
-        # Change 255 for 1
-        mask[mask==255]=1
-        
-        
-        # Back to PILMask
-        mask = PILMask.create(mask)
-        return img, mask
+repo_id = "Ignaciobfp/segmentacion-dron-marras"
+learner = from_pretrained_fastai(repo_id)
+
+device = torch.device("cpu") 
+model = learner.model
+model = model.cpu()
 
 def transform_image(image):
     my_transforms = transforms.Compose([transforms.ToTensor(),
@@ -30,20 +19,11 @@ def transform_image(image):
     image_aux = image
     return my_transforms(image_aux).unsqueeze(0).to(device)
 
-repo_id = "Ignaciobfp/segmentacion-dron-marras"
-
-learner = from_pretrained_fastai(repo_id)
-
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu") 
-model = learner.model
-model = model.cpu()
-
-
 
 # Definimos una función que se encarga de llevar a cabo las predicciones
 def predict(img):
-    #img = PILImage.create(img)
-    image = transforms.Resize((400,400))(img)
+    img_pil = PIL.Image.fromarray(img, 'RGB')
+    image = transforms.Resize((400,400))(img_pil)
     tensor = transform_image(image=image)
     model.to(device)
     with torch.no_grad():
@@ -52,7 +32,7 @@ def predict(img):
     mask = np.array(outputs.cpu())
     mask[mask==1]=255
     mask=np.reshape(mask,(400,400))
-    return mask
+    return Image.fromarray(mask.astype('uint8')
     
 # Creamos la interfaz y la lanzamos. 
 gr.Interface(fn=predict, inputs=gr.inputs.Image(shape=(400, 400)), outputs="image", examples=['examples/1CA SUR_1200_800.png', 'examples/1CA SUR_4000_1200.png', 'examples/1CA SUR_4800_2000.png']).launch(share=False)