| from huggingface_hub import from_pretrained_fastai | |
| import gradio as gr | |
| from fastai.vision.all import * | |
| import PIL | |
| import torchvision.transforms as transforms | |
| repo_id = "paascorb/practica3_Segmentation" | |
| learner = from_pretrained_fastai(repo_id) | |
| 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) | |
| def predict(img): | |
| img = PIL.Image.fromarray(img, "RGB") | |
| image = transforms.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) | |
| outputs = torch.argmax(outputs,1) | |
| mask = np.array(outputs.cpu()) | |
| mask[mask==1]=255 | |
| mask[mask==2]=150 | |
| mask[mask==3]=76 | |
| mask[mask==4]=29 | |
| mask=np.reshape(mask,(480,640)) | |
| return Image.fromarray(mask.astype('uint8')) | |
| gr.Interface(fn=predict, inputs=gr.inputs.Image(shape=(128, 128)), outputs=[gr.outputs.Image(type="pil", label="Predicci贸n")], examples=['color_155.jpg','color_154.jpg']).launch(share=False) |