import os
import cv2
import gradio as gr
import numpy as np
import spaces
import torch
import torch.nn.functional as F
from gradio.themes.utils import sizes
from torchvision import transforms
from PIL import Image
import tempfile
from classes_and_palettes import GOLIATH_PALETTE, GOLIATH_CLASSES
class Config:
ASSETS_DIR = os.path.join(os.path.dirname(__file__), 'assets')
CHECKPOINTS_DIR = os.path.join(ASSETS_DIR, "checkpoints")
CHECKPOINTS = {
"0.3b": "sapiens_0.3b_goliath_best_goliath_mIoU_7673_epoch_194_torchscript.pt2",
"0.6b": "sapiens_0.6b_goliath_best_goliath_mIoU_7777_epoch_178_torchscript.pt2",
"1b": "sapiens_1b_goliath_best_goliath_mIoU_7994_epoch_151_torchscript.pt2",
}
class ModelManager:
@staticmethod
def load_model(checkpoint_name: str):
checkpoint_path = os.path.join(Config.CHECKPOINTS_DIR, Config.CHECKPOINTS[checkpoint_name])
model = torch.jit.load(checkpoint_path)
model.eval()
model.to("cuda")
return model
@staticmethod
@torch.inference_mode()
def run_model(model, input_tensor, height, width):
output = model(input_tensor)
output = F.interpolate(output, size=(height, width), mode="bilinear", align_corners=False)
_, preds = torch.max(output, 1)
return preds
class ImageProcessor:
def __init__(self):
self.transform_fn = transforms.Compose([
transforms.Resize((1024, 768)),
transforms.ToTensor(),
transforms.Normalize(mean=[123.5/255, 116.5/255, 103.5/255], std=[58.5/255, 57.0/255, 57.5/255]),
])
@spaces.GPU
def process_image(self, image: Image.Image, model_name: str):
model = ModelManager.load_model(model_name)
input_tensor = self.transform_fn(image).unsqueeze(0).to("cuda")
preds = ModelManager.run_model(model, input_tensor, image.height, image.width)
mask = preds.squeeze(0).cpu().numpy()
# Visualize the segmentation
blended_image = self.visualize_pred_with_overlay(image, mask)
# Create downloadable .npy file
npy_path = tempfile.mktemp(suffix='.npy')
np.save(npy_path, mask)
return blended_image, npy_path
@staticmethod
def visualize_pred_with_overlay(img, sem_seg, alpha=0.5):
img_np = np.array(img.convert("RGB"))
sem_seg = np.array(sem_seg)
num_classes = len(GOLIATH_CLASSES)
ids = np.unique(sem_seg)[::-1]
legal_indices = ids < num_classes
ids = ids[legal_indices]
labels = np.array(ids, dtype=np.int64)
colors = [GOLIATH_PALETTE[label] for label in labels]
overlay = np.zeros((*sem_seg.shape, 3), dtype=np.uint8)
for label, color in zip(labels, colors):
overlay[sem_seg == label, :] = color
blended = np.uint8(img_np * (1 - alpha) + overlay * alpha)
return Image.fromarray(blended)
class GradioInterface:
def __init__(self):
self.image_processor = ImageProcessor()
def create_interface(self):
app_styles = """
"""
header_html = f"""