Update app.py

app.py

@@ -39,6 +39,37 @@ from huggingface_hub import from_pretrained_fastai
 learn = from_pretrained_fastai("islasher/segm-grapes")
 
 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)
+
+
+        # Changes: (codes= array(['Background', 'Leaves', 'Wood', 'Pole', 'Grape'], dtype='<U10'))
+        mask[mask==150]=1 #leaves
+
+        mask[mask==76]=3 #pole
+        mask[mask==74]=3 #pole
+
+        mask[mask==29]=2 #wood
+        mask[mask==25]=2 #wood
+
+        mask[mask==255]=4 #grape
+
+        mask[mask==0]=0
+        # Back to PILMask
+        mask = PILMask.create(mask)
+        return img, mask
+
+
+
 def transform_image(image):
     my_transforms = transforms.Compose([transforms.ToTensor(),
                                         transforms.Normalize(