File size: 2,635 Bytes
222641d 200fb6b 222641d 200fb6b 222641d 200fb6b 222641d |
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 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 |
import gradio as gr
import torch
import os
import torch.nn as nn
from torchvision import transforms
from sklearn.model_selection import GroupKFold
import cv2
from skimage import io
import torch
from torch import nn
import os
from datetime import datetime
import time
import random
import cv2
import pandas as pd
import numpy as np
import albumentations as A
from albumentations.pytorch.transforms import ToTensorV2
import sklearn
from efficientnet_pytorch import EfficientNet
def get_net():
net = EfficientNet.from_pretrained('efficientnet-b4')
net._fc = nn.Linear(in_features=1792, out_features=4, bias=True)
return net
def get_test_transforms(mode):
if mode == 0:
return A.Compose([
A.Resize(height=512, width=512, p=1.0),
ToTensorV2(p=1.0),
], p=1.0)
elif mode == 1:
return A.Compose([
A.HorizontalFlip(p=1),
A.Resize(height=512, width=512, p=1.0),
ToTensorV2(p=1.0),
], p=1.0)
elif mode == 2:
return A.Compose([
A.VerticalFlip(p=1),
A.Resize(height=512, width=512, p=1.0),
ToTensorV2(p=1.0),
], p=1.0)
else:
return A.Compose([
A.HorizontalFlip(p=1),
A.VerticalFlip(p=1),
A.Resize(height=512, width=512, p=1.0),
ToTensorV2(p=1.0),
], p=1.0)
def preprocess_image(image_path, transforms=None):
image = cv2.imread(image_path, cv2.IMREAD_COLOR)
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB).astype(np.float32)
image /= 255.0 # Normalize to [0, 1]
if transforms:
sample = {'image': image}
sample = transforms(**sample)
image = sample['image']
image = torch.tensor(image, dtype=torch.float32).unsqueeze(0)
return image
net = get_net()
checkpoint = torch.load('model.bin', map_location="cpu")
net.load_state_dict(checkpoint['model_state_dict']);
net.eval()
path = "uploads" # Define path for uploads
os.makedirs(path, exist_ok=True) # Ensure directory exists
def predict_image(path):
# img_path = os.path.join(path, img.name)
# img.save(img_path)
image_tensor = preprocess_image(path, get_test_transforms(224))
pred = net(image_tensor)
y_pred = 1 - nn.functional.softmax(pred, dim=1).data.cpu().numpy()[:, 0]
return float(y_pred.item())
iface = gr.Interface(
fn=predict_image,
inputs=gr.Image(type="filepath"),
outputs="number",
title="Image Prediction App",
description="Upload an image and get a model prediction."
)
iface.launch()
|