Spaces:
Running
Running
File size: 4,044 Bytes
e36eee4 | 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 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 | import sys
import os
import io
import base64
import numpy as np
from PIL import Image
import torch
from fastapi import FastAPI, File, UploadFile
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
# Add current directory to path so HF Space finds it
import sys
import os
current_dir = os.path.dirname(os.path.abspath(__file__))
if current_dir not in sys.path:
sys.path.append(current_dir)
from omegaconf import OmegaConf
from src.model import build_model
from src.attention_viz import attention_rollout_full, make_overlay
from src.dataset import QUESTION_GROUPS
from torchvision import transforms
app = FastAPI()
app.add_middleware(
CORSMiddleware,
allow_origins=["*"],
allow_credentials=True,
allow_methods=["*"],
allow_headers=["*"],
)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = None
cfg = None
transform = None
@app.on_event("startup")
def load_model():
global model, cfg, transform
print("Loading configuration...")
base_cfg = OmegaConf.load(os.path.join(current_dir, "configs/base.yaml"))
# We load the full train config
try:
exp_cfg = OmegaConf.load(os.path.join(current_dir, "configs/full_train.yaml"))
cfg = OmegaConf.merge(base_cfg, exp_cfg)
except:
cfg = base_cfg
print("Building model...")
model = build_model(cfg).to(device)
ckpt_path = os.path.join(current_dir, "best_full_train.pt")
if os.path.exists(ckpt_path):
print(f"Loading checkpoint from {ckpt_path}")
ckpt = torch.load(ckpt_path, map_location=device, weights_only=True)
model.load_state_dict(ckpt["model_state"])
else:
print(f"WARNING: Checkpoint not found at {ckpt_path}")
model.eval()
# Galaxy Zoo image transform: resize, crop, center, normalize
# Assuming standard Imagenet + ViT transforms for 224x224
transform = transforms.Compose([
transforms.Resize(224),
transforms.CenterCrop(224),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])
@app.post("/api/predict")
async def predict(file: UploadFile = File(...)):
contents = await file.read()
image = Image.open(io.BytesIO(contents)).convert("RGB")
# Transform image
img_tensor = transform(image).unsqueeze(0).to(device)
with torch.no_grad():
with torch.amp.autocast("cuda", enabled=True):
logits = model(img_tensor)
# Get attention weights
layers = model.get_all_attention_weights()
# Process predictions mapping
predictions = logits[0].cpu().numpy()
results = {}
# In proper evaluation, hierarchical softmax is applied per question group
import torch.nn.functional as F
probs = logits.detach().cpu().clone()
for q_name, (start, end) in QUESTION_GROUPS.items():
probs[:, start:end] = F.softmax(probs[:, start:end], dim=-1)
probs_np = probs[0].numpy()
for q_name, (start, end) in QUESTION_GROUPS.items():
results[q_name] = probs_np[start:end].tolist()
# Generate Attention Heatmap Overlay
if layers is not None:
# attention_rollout_full expects list of [B, H, N+1, N+1]
all_layer_attns = [l.cpu() for l in layers]
rollout_map = attention_rollout_full(all_layer_attns, patch_size=16, image_size=224)[0]
# original image numpy for overlay (denormalised size)
original_img_np = np.array(image.resize((224, 224)))
overlay = make_overlay(original_img_np, rollout_map, alpha=0.5, colormap="inferno")
# Encode to base64
overlay_img = Image.fromarray(overlay)
buffered = io.BytesIO()
overlay_img.save(buffered, format="PNG")
heatmap_base64 = base64.b64encode(buffered.getvalue()).decode('utf-8')
else:
heatmap_base64 = None
return {
"predictions": results,
"heatmap": heatmap_base64
}
|