[Python] Trying to implement CAT-Seg: Cost Aggregation for Open-Vocabulary Semantic...

So I'm trying to implement a simplified version of the above mentioned paper in Kaggle notebook, using ADE-20K dataset for training and validation. Upon testing it on a sample input, I'm getting a somewhat ok segmentation. However, the class names are shown as class 1, class 2 etc instead of their actual names. I can't seem to fix it

In the dataset file, there are text files called objectInfo.txt and sceneCategories.txt. In my code,I load them like this.

# ======================
# 2. Data Loading with Proper Class Mapping
# ======================
def load_ade20k_metadata(data_root):
"""Load class and scene metadata with robust parsing"""
classes = []
scene_map = {}

# 1. Load CLASSES from objectInfo150.txt
class_file = os.path.join(data_root, "annotations", "objectInfo150.txt")
try:
with open(class_file, 'r', encoding='utf-8') as f:
# Skip header line (Idx\tRatio\tTrain\tVal\tName)
next(f)
for line in f:
parts = line.strip().split('\t')
if len(parts) >= 5: # Ensure we have all columns
class_name = parts[4].strip().split(',')[0].strip() # Take first name before comma
classes.append(class_name)
print(f"Successfully loaded {len(classes)} classes")
print("Sample classes:", classes[:10])
except Exception as e:
print(f"Error loading classes: {e}")
classes = [f"class_{i}" for i in range(150)]

# 2. Load SCENES from sceneCategories.txt
scene_file = os.path.join(data_root, "annotations", "sceneCategories.txt")
try:
with open(scene_file, 'r', encoding='utf-8') as f:
for line in f:
if ', ' in line:
img, scene = line.strip().split(', ', 1)
scene_map = scene
scenes = sorted(set(scene_map.values()))
print(f"Loaded {len(scenes)} scene categories")
except Exception as e:
print(f"Error loading scenes: {e}")
scenes = ["scene"]
scene_map = defaultdict(lambda: "scene")

return classes, scenes, scene_map

ADE20K_CLASSES, SCENE_CATEGORIES, SCENE_MAP = load_ade20k_metadata(
"/kaggle/input/ade20k-dataset/ADEChallengeData2016"
)
config.num_classes = len(ADE20K_CLASSES)
config.num_scenes = len(SCENE_CATEGORIES)

[url="https://stackoverflow.com/questions/79712596/trying-to-implement-cat-seg-cost-aggregation-for-open-vocabulary-semantic-segme"]Continue reading...[/url]