Spaces:
Running
Running
#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) |