Spaces:
Running
Running
import gradio as gr | |
from glob import glob | |
import os | |
import time | |
from PIL import Image | |
import albumentations as A | |
import cv2 | |
import numpy as np | |
import matplotlib.patches as mpatches | |
import matplotlib.pyplot as plt | |
import pandas as pd | |
from scipy.ndimage.morphology import binary_dilation | |
import segmentation_models_pytorch as smp | |
from sklearn.impute import SimpleImputer | |
from sklearn.model_selection import train_test_split | |
import torch | |
import torch.nn as nn | |
from torch.optim import Adam | |
from torch.optim.lr_scheduler import ReduceLROnPlateau | |
from torch.utils.data import Dataset, DataLoader | |
from torchvision import transforms as T | |
from tqdm import tqdm | |
from tensorflow.keras.models import load_model | |
model = smp.MAnet( | |
encoder_name="efficientnet-b7", | |
encoder_weights="imagenet", | |
in_channels=3, | |
classes=1, | |
activation='sigmoid',) | |
transform = A.Compose([ | |
A.ChannelDropout(p=0.3), | |
A.RandomBrightnessContrast(p=0.3), | |
A.ColorJitter(p=0.3), | |
]) | |
model.load_state_dict(torch.load("weights.pt", map_location=torch.device('cpu'))) | |
model.eval() | |
def segment(image): | |
image = transform(image=image) | |
image = image.get("image") | |
image = T.functional.to_tensor(image) | |
prediction = model(image[None, ...]) | |
prediction = np.squeeze(prediction.detach().numpy()) | |
return Image.fromarray(prediction) | |
iface = gr.Interface(fn=segment, inputs="image", outputs="image").launch() |