Spaces:
Running
Running
File size: 1,407 Bytes
44a9d05 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 |
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() |