app.py

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()