Spaces:
Sleeping
Sleeping
| from huggingface_hub import from_pretrained_fastai | |
| import gradio as gr | |
| 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 | |
| def transform_image(image): | |
| my_transforms = transforms.Compose([transforms.ToTensor(), | |
| transforms.Normalize( | |
| [0.485, 0.456, 0.406], | |
| [0.229, 0.224, 0.225])]) | |
| 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) | |
| tensor = transform_image(image=image) | |
| model.to(device) | |
| with torch.no_grad(): | |
| outputs = model(tensor) | |
| outputs = torch.argmax(outputs,1) | |
| mask = np.array(outputs.cpu()) | |
| mask[mask==1]=255 | |
| mask=np.reshape(mask,(400,400)) | |
| return mask | |
| # 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) | |