File size: 1,696 Bytes
8d1d038 a2fa090 863a1f1 a2fa090 5570c58 2933508 8d1d038 2933508 8d1d038 ea5ffee 8d1d038 5570c58 f462b26 8e1e069 5570c58 f462b26 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 |
from huggingface_hub import from_pretrained_fastai
import gradio as gr
from fastai.vision.all import *
from albumentations import (
Compose,
OneOf,
ElasticTransform,
GridDistortion,
OpticalDistortion,
HorizontalFlip,
VerticalFlip,
Rotate,
Transpose,
CLAHE,
ShiftScaleRotate,
RandomScale
)
def get_y_fn (x):
return Path(str(x).replace("Images","Labels").replace("color","gt").replace(".jpg",".png"))
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
mask = np.array(mask)
new_mask = np.zeros_like(mask, dtype=np.uint8)
new_mask[mask==150]=1 #Clase Leaves
new_mask[(mask==29) | (mask==25)]=2 #Clase Wood
new_mask[(mask==74) | (mask==76)]=3 #Clase Pole
new_mask[mask==255]=4 #Clase Wood
mask = PILMask.create(new_mask)
return img, mask
repo_id = "joortif/unet-resnet34-segmentation"
learn = from_pretrained_fastai(repo_id)
def segment_image(img):
pred, _, _ = learn.predict(img)
return pred
interface = gr.Interface(
fn=segment_image,
inputs=gr.Image(),
outputs=gr.Image(),
examples=["color_184.jpg","color_154.jpg","color_180.jpg"],
title="Segmentación Semántica con FastAI",
).launch(share=False) |