Spaces:
Sleeping
Sleeping
| #app.py: | |
| # from huggingface_hub import from_pretrained_fastai | |
| import gradio as gr | |
| from fastai import * | |
| from fastai.data.block import DataBlock | |
| from fastai.data.transforms import get_image_files, FuncSplitter, Normalize | |
| from fastai.layers import Mish | |
| from fastai.losses import BaseLoss | |
| from fastai.optimizer import ranger | |
| from fastai.torch_core import tensor | |
| from fastai.vision.augment import aug_transforms | |
| from fastai.vision.core import PILImage, PILMask | |
| from fastai.vision.data import ImageBlock, MaskBlock, imagenet_stats | |
| from fastai.vision.learner import unet_learner | |
| from PIL import Image | |
| import numpy as np | |
| from torch import nn | |
| import torch | |
| import torch.nn.functional as F | |
| #from __future__ import annotations | |
| #from nbdev.showdoc import * | |
| #from fastai import fastcore | |
| from fastcore.test import * | |
| from fastcore.nb_imports import * | |
| from fastcore.imports import * | |
| from fastcore.foundation import * | |
| from fastcore.utils import * | |
| from fastcore.dispatch import * | |
| from fastcore.transform import * | |
| import inspect | |
| # # repo_id = "YOUR_USERNAME/YOUR_LEARNER_NAME" | |
| # repo_id = "islasher/segm-grapes" | |
| # repo_id='islasher/segm-grapes' | |
| # # Definimos una funci贸n que se encarga de llevar a cabo las predicciones | |
| # from fastai.learner import load_learner | |
| # # Cargar el modelo y el tokenizador | |
| # learn = load_learner(repo_id) | |
| #learner = from_pretrained_fastai(repo_id) | |
| class ItemTransform(Transform): | |
| "A transform that always take tuples as items" | |
| _retain = True | |
| def __call__(self, x, **kwargs): return self._call1(x, '__call__', **kwargs) | |
| def decode(self, x, **kwargs): return self._call1(x, 'decode', **kwargs) | |
| def _call1(self, x, name, **kwargs): | |
| if not _is_tuple(x): return getattr(super(), name)(x, **kwargs) | |
| y = getattr(super(), name)(list(x), **kwargs) | |
| if not self._retain: return y | |
| if is_listy(y) and not isinstance(y, tuple): y = tuple(y) | |
| return retain_type(y, x) | |
| from huggingface_hub import from_pretrained_fastai | |
| import torchvision.transforms as transforms | |
| # from Transform import ItemTransform | |
| from albumentations import ( | |
| Compose, | |
| OneOf, | |
| ElasticTransform, | |
| GridDistortion, | |
| OpticalDistortion, | |
| HorizontalFlip, | |
| Rotate, | |
| Transpose, | |
| CLAHE, | |
| ShiftScaleRotate | |
| ) | |
| 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 | |
| #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 get_y_fn (x): | |
| return Path(str(x).replace("Images","Labels").replace("color","gt").replace(".jpg",".png")) | |
| learn = from_pretrained_fastai("islasher/segm-grapes") | |
| #device = torch.device("cuda" if torch.cuda.is_available() else "cpu") | |
| 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) | |
| # Definimos una funci贸n que se encarga de llevar a cabo las predicciones | |
| def predict(img): | |
| img=Image.fromarray(img) | |
| image = transforms.Resize((480,640))(img) | |
| tensor = transform_image(image=image) | |
| with torch.no_grad(): | |
| outputs = learn.model(tensor) | |
| outputs = torch.argmax(outputs,1) | |
| mask = np.array(outputs) | |
| mask[mask==1]=150 | |
| mask[mask==3]=76 #pole # y no 74 | |
| # mask[mask==5]=74 #pole | |
| mask[mask==2]=29 #wood # y no 25 | |
| # mask[mask==6]=25 #wood | |
| mask[mask==4]=255 #grape | |
| mask=np.reshape(mask,(480,640)) #en modo matriz | |
| return Image.fromarray(mask.astype('uint8')) | |
| # Creamos la interfaz y la lanzamos. | |
| gr.Interface(fn=predict, inputs=gr.Image(), outputs=gr.Image(),examples=['color_154.jpg','color_155.jpg']).launch(share=False) #shape=(128, 128) shape=(480,640) |