kalpkanungo/SceneGraphNet

1import torch2import cv23import json4import os5from src.model import RelationshipNet6from huggingface_hub import hf_hub_download7 8 9MODEL_REPO = "kalpkanungo/scenegraphnet-relationship-model"10MODEL_FILENAME = "relationship_model.pth"11 12LABEL_MAP_PATH = "data/relationship_dataset/label_map.json"13 14device = "cuda" if torch.cuda.is_available() else "cpu"15 16 17if os.path.exists(LABEL_MAP_PATH):18    with open(LABEL_MAP_PATH) as f:19        label_map = json.load(f)20else:21    print("⚠️ label_map.json not found, using fallback")22    label_map = {23        "0": "on",24        "1": "next to",25        "2": "under"26    }27 28inv_map = {v: k for k, v in label_map.items()}29num_classes = len(label_map)30 31 32model = RelationshipNet(num_classes)33 34try:35    model_path = hf_hub_download(36        repo_id=MODEL_REPO,37        filename=MODEL_FILENAME38    )39    print("✅ Model downloaded from Hugging Face")40 41    model.load_state_dict(torch.load(model_path, map_location=device))42    model.to(device)43    model.eval()44 45except Exception as e:46    print(f"⚠️ Failed to load model from HF: {e}")47    model = None48 49 50def predict(image):51   52    if model is None:53        return "next to"54 55    image = cv2.resize(image, (128, 128))56    image = image / 255.057    image = (image - 0.5) / 0.558 59    image = torch.tensor(image, dtype=torch.float32).permute(2, 0, 1)60    image = image.unsqueeze(0).to(device)61 62    with torch.no_grad():63        output = model(image)64        pred = torch.argmax(output, dim=1).item()65 66    return inv_map.get(pred, "unknown")