tesa.py

import logging

import torch
import torch.nn as nn
import torch.nn.functional as F
from mmengine.structures import PixelData
from mmseg.models.data_preprocessor import SegDataPreProcessor
from mmseg.models.segmentors import BaseSegmentor
from mmseg.registry import MODELS

import clip
from pamr import PAMR
from prompts.imagenet_template import openai_imagenet_template, get_prompt_bank, detect_dataset_from_path

# ============================================================================
# CTE: Contrastive Text Enhancement
# ============================================================================
# Use contrastive descriptions to improve class discrimination in CLIP space.
#
# Typical confusions include terrain/vegetation, wall/building, and
# rider/person.
#
# Each class receives descriptions that emphasize differences from its peers.

SEMANTIC_RELATIONS = {
    # Format: class name -> contrastive descriptions.

    # ============ Cityscapes classes ============
    "terrain": [
        "bare ground without any plants or grass",
        "brown dirt or soil surface, not green vegetation",
        "unpaved earth ground, distinct from grass or trees",
        "muddy or sandy ground with no plant life",
    ],
    "vegetation": [
        "green plants, trees, bushes with leaves",
        "living plants and foliage, not bare ground",
        "green leafy vegetation, distinct from brown terrain",
    ],
    "wall": [
        "a standalone wall structure, not part of a building facade",
        "a low barrier wall separate from buildings",
        "a concrete or brick wall that is not a building",
        "a retaining wall or boundary wall, distinct from building walls",
    ],
    "building": [
        "a multi-story building structure with windows",
        "an architectural building, not a simple wall",
        "a complete building facade, not just a wall segment",
    ],
    "rider": [
        "a person actively riding on a bicycle or motorcycle",
        "someone sitting on and controlling a two-wheeled vehicle",
        "a cyclist or motorcyclist, not a pedestrian walking",
    ],
    "person": [
        "a pedestrian walking on foot, not riding anything",
        "a standing or walking person, not on a vehicle",
    ],
    "sidewalk": [
        "a paved pedestrian path beside the road",
        "a walkway for people, distinct from the vehicle road",
        "concrete pavement for walking, not for driving",
    ],
    "road": [
        "a paved surface for vehicles to drive on",
        "the main driving lane, not the pedestrian sidewalk",
        "asphalt road for cars, distinct from sidewalk",
    ],

    # ============ VOC20 classes ============
    # Confusion group: cat / dog (furry quadruped pets)
    "cat": [
        "a feline animal with a small face and pointed ears",
        "a cat, not a dog, with a slender body and whiskers",
        "a domestic cat distinct from a dog by its retractable claws",
    ],
    "dog": [
        "a canine animal with a snout and floppy or pointed ears",
        "a dog, not a cat, with a loyal and expressive face",
        "a domestic dog distinct from a cat by its larger muzzle",
    ],
    # Confusion group: horse / cow / sheep (farm animals)
    "horse": [
        "a tall equine animal with a long mane and slender legs",
        "a horse, not a cow or sheep, with a muscular athletic body",
        "a horse with hooves and a flowing tail, distinct from cattle",
    ],
    "cow": [
        "a stocky bovine animal with a broad body and udder",
        "a cow, not a horse or sheep, with a wide face and horns",
        "a cattle animal with black-and-white or brown hide",
    ],
    "sheep": [
        "a woolly animal covered in thick fleece",
        "a sheep, not a cow or horse, with a compact fluffy body",
        "a sheep with curly wool coat, distinct from smooth-skinned cattle",
    ],
    # Confusion group: chair / sofa / table (furniture)
    "chair": [
        "a single-seat furniture piece with a backrest",
        "a chair, not a sofa or table, designed for one person to sit",
        "a small rigid seat distinct from a cushioned couch",
    ],
    "sofa": [
        "a large cushioned seating furniture for multiple people",
        "a sofa or couch, not a single chair, with soft upholstery",
        "a wide padded couch distinct from a rigid single chair",
    ],
    "table": [
        "a flat-topped furniture piece supported by legs for placing objects",
        "a table, not a chair or sofa, with a horizontal surface",
        "a dining or work table distinct from seating furniture",
    ],
    # Confusion group: bicycle / motorbike (two-wheeled vehicles)
    "bicycle": [
        "a human-powered two-wheeled vehicle with pedals",
        "a bicycle, not a motorbike, propelled by pedaling",
        "a lightweight pedal bike without an engine",
    ],
    "motorbike": [
        "a motorized two-wheeled vehicle with an engine",
        "a motorbike, not a bicycle, powered by a motor",
        "a heavy motorcycle with engine, distinct from a pedal bicycle",
    ],
    # Confusion group: bus / car / train (vehicles)
    "bus": [
        "a large multi-passenger vehicle for public transit",
        "a bus, not a car or train, with a long rectangular body",
        "a big road vehicle carrying many passengers, distinct from small cars",
    ],
    "car": [
        "a small four-wheeled passenger automobile",
        "a car, not a bus or train, compact enough for a few passengers",
        "a personal motor vehicle distinct from large public transport",
    ],
    "train": [
        "a rail vehicle running on tracks with multiple carriages",
        "a train, not a bus or car, moving along railway tracks",
        "a long locomotive on rails, distinct from road vehicles",
    ],
    # Confusion group: bird / aeroplane (flying objects)
    "bird": [
        "a living animal with feathers and wings that can fly",
        "a bird, not an airplane, a small organic creature",
        "a feathered bird distinct from a metallic aircraft",
    ],
    "aeroplane": [
        "a large metallic aircraft with fixed wings and jet engines",
        "an airplane, not a bird, a man-made flying machine",
        "a mechanical aircraft distinct from a living bird",
    ],
    # ship (water vs land)
    "ship": [
        "a large vessel floating on water for transportation",
        "a ship or boat on water, not a land vehicle",
        "a watercraft distinct from cars, buses, and trains",
    ],
    # Confusion group: bottle / pottedplant (small objects with similar silhouettes)
    "bottle": [
        "a narrow cylindrical container for holding liquids",
        "a bottle, not a plant, a smooth manufactured object",
        "a glass or plastic bottle distinct from a living potted plant",
    ],
    "pottedplant": [
        "a living plant growing in a pot or container",
        "a potted plant, not a bottle, with green leaves and stems",
        "a houseplant in a container distinct from a manufactured bottle",
    ],
    # television monitor
    "television monitor": [
        "a flat electronic display screen for showing images",
        "a TV or monitor with a glowing screen surface",
        "an electronic screen device distinct from furniture",
    ],

    # ============ VOC21 background scene elements ============
    # These are aliases of the background class; CTE helps distinguish them
    # from visually similar foreground classes.

    # desk vs table (foreground) — most confusable pair
    "desk": [
        "a work desk or study desk, typically in an office or home",
        "a desk for writing or working, not a dining table",
        "an office desk with drawers, distinct from a dining table",
    ],
    # bed vs sofa (foreground)
    "bed": [
        "a sleeping bed with mattress and pillows in a bedroom",
        "a bed for sleeping, not a sofa for sitting",
        "a bed frame with bedding, distinct from a living room couch",
    ],
    # tree vs pottedplant (foreground)
    "tree": [
        "a large outdoor tree with a trunk and canopy of branches",
        "a tall tree rooted in the ground, not a small potted plant",
        "a tree growing outdoors, distinct from a houseplant in a pot",
    ],
    # sky vs bird/aeroplane (foreground flying objects)
    "sky": [
        "the open sky above, a vast background area",
        "sky without objects, distinct from birds or airplanes in it",
        "blue or cloudy sky as the scene background",
    ],
    # grass vs green foreground objects
    "grass": [
        "green grass covering the ground surface",
        "a grassy lawn area, not a potted plant or tree",
        "ground-level grass, distinct from tall trees or bushes",
    ],
    # sea vs ship (foreground watercraft)
    "sea": [
        "a vast body of ocean water as scene background",
        "sea water surface, not a ship or boat floating on it",
        "the ocean or sea, distinct from watercraft vessels",
    ],
    # river (water body, distinct from sea)
    "river": [
        "a flowing freshwater river or stream",
        "a river, not the open sea, a narrow body of water inland",
        "flowing river water, distinct from ocean waves",
    ],
    # mountain vs hill
    "mountain": [
        "a tall rocky mountain with steep slopes and peaks",
        "a mountain, not a small hill, with significant elevation",
        "a large mountain in the background landscape",
    ],
    "hill": [
        "a gentle low-elevation hill or mound",
        "a rolling hill, not a steep mountain, with gradual slopes",
        "a small grassy hill distinct from a tall rocky mountain",
    ],
    # door vs window
    "door": [
        "an entrance door in a wall for people to walk through",
        "a door, not a window, opaque and full-height for entry",
        "a door opening distinct from a glass window",
    ],
    "window": [
        "a glass window in a wall for letting in light",
        "a window, not a door, with glass panes for seeing through",
        "a window opening distinct from an entrance door",
    ],
    # floor vs road (already has entry)
    "floor": [
        "an indoor floor surface of a room",
        "a floor inside a building, not an outdoor road",
        "tiled or wooden floor, distinct from outdoor pavement",
    ],
    # fence vs wall (already has entry)
    "fence": [
        "an open or slatted barrier around a property",
        "a fence, not a solid wall, with gaps between slats or wires",
        "a lightweight fence barrier distinct from a solid concrete wall",
    ],

    # ============ Context59 additional classes ============
    # Ground/terrain surfaces (highly confusable group)
    "ground": [
        "bare outdoor ground or earth surface",
        "ground, not grass or road, unpaved dirt or soil",
        "exposed ground surface distinct from paved areas",
    ],
    "rock": [
        "a natural rock or stone formation",
        "rock, not ground or mountain, a solid stone surface",
        "rocky terrain distinct from bare dirt ground",
    ],
    "snow": [
        "white snow covering the ground or objects",
        "snow, not water or cloud, frozen white precipitation",
        "snowy surface distinct from other ground cover",
    ],
    "water": [
        "a body of water such as a lake, pond, or pool",
        "water surface, not snow or sky, a liquid body",
        "water in a natural or artificial setting",
    ],
    "platform": [
        "a raised flat surface for standing or boarding",
        "a platform, not a floor or sidewalk, elevated above ground",
        "a station platform or stage distinct from regular ground",
    ],
    "track": [
        "railway tracks or a running track on the ground",
        "a track, not a road, with rails or lane markings",
        "railroad tracks or athletic track distinct from a road",
    ],
    # Seating (bench vs chair vs sofa)
    "bench": [
        "a long seat for multiple people, typically outdoors",
        "a bench, not a chair, a wider seat without arms",
        "a park or public bench distinct from an indoor chair",
    ],
    # Fabric (cloth vs bedclothes vs curtain)
    "cloth": [
        "a piece of fabric or textile material",
        "cloth, not curtain or bedding, a loose fabric",
        "fabric material distinct from structured textiles",
    ],
    "bedclothes": [
        "bed sheets, blankets, and pillows on a bed",
        "bedclothes, not curtain or loose cloth, bedding fabric",
        "bed linens and covers distinct from other fabrics",
    ],
    # Electronics (computer vs tvmonitor vs keyboard vs mouse)
    "computer": [
        "a computer tower or laptop device",
        "a computer, not a TV monitor, a computing machine",
        "a PC or laptop distinct from a display screen",
    ],
    "tvmonitor": [
        "a flat electronic display screen for showing images",
        "a TV or monitor screen, not a computer tower",
        "a display screen distinct from a computing device",
    ],
    # Vehicles (truck vs bus)
    "truck": [
        "a cargo or freight truck for hauling goods",
        "a truck, not a bus, designed for carrying cargo",
        "a heavy truck distinct from a passenger bus",
    ],
    "boat": [
        "a small watercraft floating on water",
        "a boat, not a car or bus, a vessel on water",
        "a watercraft distinct from land vehicles",
    ],
    # Storage (cabinet vs shelves)
    "cabinet": [
        "a closed storage furniture with doors",
        "a cabinet, not open shelves, enclosed storage",
        "a closed cabinet distinct from open shelving",
    ],
    "shelves": [
        "open horizontal shelving for storage or display",
        "shelves, not a closed cabinet, open racks",
        "open shelves distinct from enclosed cabinets",
    ],
    # Plants (flower vs pottedplant vs tree)
    "flower": [
        "a blooming flower with colorful petals",
        "a flower, not a potted plant or tree, just the blossom",
        "flowers with petals distinct from green foliage plants",
    ],
    # Wood vs tree
    "wood": [
        "wooden material or surface, processed timber",
        "wood, not a living tree, cut or processed lumber",
        "wooden surface distinct from a living tree with leaves",
    ],
    # Ceiling vs sky vs floor
    "ceiling": [
        "the overhead surface inside a room",
        "a ceiling, not the sky, an indoor overhead surface",
        "an indoor ceiling distinct from the open sky above",
    ],
    # Sign
    "sign": [
        "a posted sign with text or symbols for information",
        "a sign, not a window or monitor, a flat board with writing",
        "a signboard distinct from electronic displays",
    ],
    # Light
    "light": [
        "a light fixture or lamp for illumination",
        "a light source, not a window, an artificial lamp",
        "a light fixture distinct from natural sunlight",
    ],

    # ============ ADE20K additional classes (150 classes) ============
    # Seating furniture (armchair vs chair vs sofa vs seat vs stool)
    "armchair": [
        "a comfortable chair with armrests",
        "an armchair, not a regular chair, with padded arms",
        "a cushioned armchair distinct from a simple chair",
    ],
    "seat": [
        "a generic seating surface or place to sit",
        "a seat, not a specific chair or sofa, a sitting area",
        "a seat surface distinct from a full chair",
    ],
    "stool": [
        "a backless seat typically at counter height",
        "a stool, not a chair, without a backrest",
        "a bar stool or kitchen stool distinct from a chair",
    ],
    "ottoman": [
        "a cushioned footstool or low seat",
        "an ottoman, not a chair or sofa, a padded footrest",
        "a low ottoman distinct from seating furniture",
    ],
    "swivelchair": [
        "a rotating office chair on wheels",
        "a swivel chair, not a fixed chair, that rotates",
        "an office swivel chair distinct from a regular chair",
    ],
    # Tables (table vs desk vs coffeetable vs countertop vs counter)
    "coffeetable": [
        "a low table in a living room for drinks and magazines",
        "a coffee table, not a dining table, low and small",
        "a living room coffee table distinct from a work desk",
    ],
    "countertop": [
        "a flat work surface in a kitchen or bathroom",
        "a countertop, not a table, a built-in work surface",
        "a kitchen countertop distinct from a freestanding table",
    ],
    "counter": [
        "a service counter or bar counter",
        "a counter, not a table, for serving or transactions",
        "a shop counter distinct from a dining table",
    ],
    "pooltable": [
        "a billiard or pool table with green felt",
        "a pool table, not a regular table, for playing billiards",
        "a game table with pockets distinct from a dining table",
    ],
    # Storage (wardrobe vs cabinet vs bookcase vs chestofdrawers vs buffet)
    "wardrobe": [
        "a tall cabinet for storing clothes",
        "a wardrobe, not a cabinet, specifically for clothing",
        "a clothes wardrobe distinct from a kitchen cabinet",
    ],
    "bookcase": [
        "a tall shelf unit for storing books",
        "a bookcase, not a cabinet, with open shelves for books",
        "a book shelf distinct from a closed cabinet",
    ],
    "chestofdrawers": [
        "a furniture piece with multiple drawers for storage",
        "a chest of drawers, not a cabinet, with pull-out drawers",
        "a dresser with drawers distinct from a wardrobe",
    ],
    "buffet": [
        "a sideboard or serving table in a dining room",
        "a buffet, not a cabinet, for serving food",
        "a dining room buffet distinct from a kitchen cabinet",
    ],
    # Lighting (lamp vs chandelier vs sconce vs streetlight)
    "lamp": [
        "a portable light fixture with a shade",
        "a lamp, not a ceiling light, a table or floor lamp",
        "a lamp with shade distinct from overhead lighting",
    ],
    "chandelier": [
        "an ornate hanging light fixture with multiple lights",
        "a chandelier, not a simple lamp, decorative ceiling light",
        "a hanging chandelier distinct from a table lamp",
    ],
    "sconce": [
        "a wall-mounted light fixture",
        "a sconce, not a lamp, attached to the wall",
        "a wall sconce distinct from a ceiling light",
    ],
    "streetlight": [
        "a tall outdoor light pole for street illumination",
        "a streetlight, not an indoor lamp, for outdoor lighting",
        "a street lamp post distinct from indoor lighting",
    ],
    # Water features (river vs lake vs sea vs swimmingpool vs fountain vs waterfall)
    "lake": [
        "a large inland body of still water",
        "a lake, not a river, a calm body of water",
        "a lake distinct from flowing river water",
    ],
    "swimmingpool": [
        "an artificial pool for swimming",
        "a swimming pool, not a natural lake, a man-made pool",
        "a pool with clear water distinct from natural water bodies",
    ],
    "fountain": [
        "a decorative water feature with jets",
        "a fountain, not a pool, with water spraying upward",
        "a water fountain distinct from a still pool",
    ],
    "waterfall": [
        "water cascading down from a height",
        "a waterfall, not a river, water falling vertically",
        "a waterfall distinct from flowing river water",
    ],
    # Screens (screen vs crtscreen vs monitor vs televisionreceiver)
    "screen": [
        "a flat display surface or projection screen",
        "a screen, not a monitor, a viewing surface",
        "a projection screen distinct from a TV",
    ],
    "crtscreen": [
        "an old-style cathode ray tube display",
        "a CRT screen, not a flat screen, a bulky old monitor",
        "a CRT display distinct from a flat panel",
    ],
    "monitor": [
        "a computer display screen",
        "a monitor, not a TV, for computer use",
        "a computer monitor distinct from a television",
    ],
    "televisionreceiver": [
        "a television set for watching broadcasts",
        "a TV, not a computer monitor, for entertainment",
        "a television distinct from a computer monitor",
    ],
    # Outdoor structures (bridge vs pier vs tower vs skyscraper)
    "pier": [
        "a platform extending into water for boats",
        "a pier, not a bridge, a dock structure",
        "a boat pier distinct from a bridge",
    ],
    "tower": [
        "a tall narrow structure or building",
        "a tower, not a regular building, tall and narrow",
        "a tower structure distinct from a wide building",
    ],
    "skyscraper": [
        "a very tall multi-story building",
        "a skyscraper, not a regular building, extremely tall",
        "a high-rise skyscraper distinct from low buildings",
    ],
    # Ground surfaces (earth vs land vs field vs path vs dirttrack)
    "earth": [
        "bare soil or dirt ground",
        "earth, not grass or rock, exposed soil",
        "bare earth distinct from paved surfaces",
    ],
    "land": [
        "a general area of ground or terrain",
        "land, not water, solid ground area",
        "land surface distinct from water bodies",
    ],
    "field": [
        "an open area of grass or crops",
        "a field, not a forest, an open grassy area",
        "an agricultural or sports field distinct from a lawn",
    ],
    "path": [
        "a narrow walkway or trail",
        "a path, not a road, a narrow walking route",
        "a footpath distinct from a wide road",
    ],
    "dirttrack": [
        "an unpaved road or racing track",
        "a dirt track, not a paved road, an unpaved surface",
        "a dirt road distinct from asphalt",
    ],
    # Plants (plant vs palm vs tree)
    "plant": [
        "a small plant or shrub",
        "a plant, not a tree, a smaller vegetation",
        "a small plant distinct from a tall tree",
    ],
    "palm": [
        "a palm tree with fronds",
        "a palm, not a regular tree, with fan-like leaves",
        "a palm tree distinct from deciduous trees",
    ],
    # Bathroom (bathtub vs sink vs toilet vs shower)
    "bathtub": [
        "a large tub for bathing",
        "a bathtub, not a sink, for full-body bathing",
        "a bath tub distinct from a shower",
    ],
    "shower": [
        "a shower stall or enclosure",
        "a shower, not a bathtub, for standing wash",
        "a shower area distinct from a bathtub",
    ],
    # Kitchen appliances (stove vs oven vs microwave vs refrigerator vs dishwasher vs washer)
    "stove": [
        "a cooking appliance with burners",
        "a stove, not an oven, with cooktop burners",
        "a stovetop distinct from an oven",
    ],
    "oven": [
        "an enclosed cooking appliance for baking",
        "an oven, not a stove, an enclosed baking chamber",
        "a baking oven distinct from a stovetop",
    ],
    "microwave": [
        "a small appliance for quick heating",
        "a microwave, not an oven, for rapid heating",
        "a microwave oven distinct from a regular oven",
    ],
    "refrigerator": [
        "a large appliance for keeping food cold",
        "a refrigerator, not a cabinet, for food storage",
        "a fridge distinct from a freezer or cabinet",
    ],
    "dishwasher": [
        "an appliance for washing dishes",
        "a dishwasher, not a washer, for cleaning dishes",
        "a dish washing machine distinct from a clothes washer",
    ],
    "washer": [
        "a clothes washing machine",
        "a washer, not a dishwasher, for laundry",
        "a laundry washer distinct from a dishwasher",
    ],
    # Misc furniture (rug vs blanket vs cushion vs pillow)
    "rug": [
        "a floor covering smaller than a carpet",
        "a rug, not a carpet, a smaller floor mat",
        "an area rug distinct from wall-to-wall carpet",
    ],
    "blanket": [
        "a soft covering for warmth",
        "a blanket, not a rug, for covering a person",
        "a bed blanket distinct from a floor rug",
    ],
    "cushion": [
        "a soft pad for sitting or leaning",
        "a cushion, not a pillow, for seating comfort",
        "a seat cushion distinct from a bed pillow",
    ],
    "pillow": [
        "a soft support for the head during sleep",
        "a pillow, not a cushion, for sleeping",
        "a bed pillow distinct from a decorative cushion",
    ],
    # Vehicles (van vs bus vs truck vs airplane vs minibike)
    "van": [
        "a large enclosed vehicle for cargo or passengers",
        "a van, not a car, a larger enclosed vehicle",
        "a cargo van distinct from a sedan",
    ],
    "airplane": [
        "a fixed-wing aircraft for flying",
        "an airplane, not a bird, a flying machine",
        "an aircraft distinct from a helicopter",
    ],
    "minibike": [
        "a small motorcycle or motorbike",
        "a minibike, not a bicycle, a small motorized bike",
        "a small motorbike distinct from a full motorcycle",
    ],
    # ============ COCO-Object specific entries ============
    "motorcycle": [
        "a motorized two-wheeled vehicle with engine, not a bicycle",
        "a motorcycle, larger than a bicycle, with motor power",
        "a motorbike distinct from a pedal bicycle",
    ],
    "couch": [
        "a large cushioned sofa for sitting, not a single chair",
        "a couch or sofa, wider than a chair, for multiple people",
        "a living room couch distinct from a bed or chair",
    ],
    "potted plant": [
        "a plant growing in a pot or container, not a tree",
        "a small indoor potted plant, not outdoor vegetation",
        "a houseplant in a pot distinct from garden plants",
    ],
    "dining table": [
        "a table used for eating meals, not a desk or counter",
        "a dining table with flat surface, distinct from a desk",
        "a table for dining, often with chairs around it",
    ],
    "tv": [
        "a television screen or monitor for viewing, not a laptop",
        "a TV display mounted or standing, distinct from a computer monitor",
        "a television set, larger than a laptop screen",
    ],
    "laptop": [
        "a portable computer with attached screen and keyboard",
        "a laptop computer, not a desktop monitor or TV",
        "a folding laptop distinct from a desktop or TV",
    ],
    "mouse": [
        "a computer mouse input device, not an animal",
        "a small handheld pointing device for computers",
        "a computer mouse with buttons, distinct from a remote",
    ],
    "remote": [
        "a handheld remote control for TV or devices",
        "a remote control with buttons, not a cell phone",
        "a small remote device distinct from a phone or mouse",
    ],
    "keyboard": [
        "a computer keyboard with keys for typing",
        "a keyboard input device, not a piano keyboard",
        "a typing keyboard distinct from a laptop",
    ],
    "cell phone": [
        "a handheld mobile phone or smartphone",
        "a cell phone with screen, not a remote control",
        "a smartphone device distinct from a remote or mouse",
    ],
    "cup": [
        "a small drinking vessel, not a bowl or wine glass",
        "a cup for coffee or tea, distinct from a bowl",
        "a drinking cup with or without handle",
    ],
    "wine glass": [
        "a stemmed glass for drinking wine",
        "a wine glass with stem, not a regular cup or bowl",
        "a tall stemmed glass distinct from a cup",
    ],
    "fork": [
        "a pronged utensil for eating, not a knife or spoon",
        "a fork with tines for picking up food",
        "a dining fork distinct from a knife or spoon",
    ],
    "knife": [
        "a cutting utensil with a blade, not a fork or spoon",
        "a knife for cutting food, distinct from scissors",
        "a table knife or kitchen knife with sharp edge",
    ],
    "spoon": [
        "a rounded utensil for scooping, not a fork or knife",
        "a spoon with concave bowl for eating",
        "a dining spoon distinct from a fork",
    ],
    "bowl": [
        "a round deep dish for food, not a cup or plate",
        "a bowl container wider than a cup",
        "a food bowl distinct from a vase or cup",
    ],
    "banana": [
        "a curved yellow fruit, not an orange or apple",
        "a banana fruit with peel, elongated shape",
        "a banana distinct from other round fruits",
    ],
    "apple": [
        "a round fruit, red or green, not an orange",
        "an apple fruit distinct from an orange by color and texture",
        "a smooth round apple, not a citrus fruit",
    ],
    "orange": [
        "a round citrus fruit with orange skin",
        "an orange fruit, not an apple, with textured peel",
        "a citrus orange distinct from an apple or sports ball",
    ],
    "broccoli": [
        "a green vegetable with tree-like florets",
        "broccoli with green crown, not a potted plant",
        "a broccoli vegetable distinct from other greens",
    ],
    "carrot": [
        "an orange root vegetable, elongated shape",
        "a carrot, not a banana, orange and tapered",
        "a carrot vegetable distinct from other foods",
    ],
    "hot dog": [
        "a sausage in a bun, not a sandwich or pizza",
        "a hot dog with bread roll, distinct from a sandwich",
        "a hot dog food item with elongated shape",
    ],
    "pizza": [
        "a flat round food with toppings and cheese",
        "a pizza pie, not a cake or sandwich",
        "a pizza with crust and toppings, distinct from cake",
    ],
    "donut": [
        "a ring-shaped fried pastry with glaze",
        "a donut with hole in center, not a cake",
        "a sweet donut distinct from a cake or pizza",
    ],
    "cake": [
        "a sweet baked dessert with frosting, not a pizza",
        "a cake, layered or frosted, distinct from a donut",
        "a birthday or dessert cake, not bread",
    ],
    "sandwich": [
        "bread with filling between slices, not a hot dog",
        "a sandwich with layers, distinct from pizza",
        "a food sandwich with bread and filling",
    ],
    "backpack": [
        "a bag worn on the back with straps",
        "a backpack, not a handbag or suitcase",
        "a carrying backpack distinct from a handbag",
    ],
    "handbag": [
        "a small bag carried by hand or shoulder",
        "a handbag or purse, not a backpack or suitcase",
        "a fashion handbag distinct from a backpack",
    ],
    "suitcase": [
        "a large travel luggage with handle",
        "a suitcase for travel, not a backpack or handbag",
        "a wheeled suitcase distinct from a bag",
    ],
    "umbrella": [
        "a canopy on a stick for rain or shade",
        "an umbrella, not a kite, for weather protection",
        "an open or closed umbrella distinct from a parasol",
    ],
    "tie": [
        "a necktie worn around the collar",
        "a formal tie, not a scarf or belt",
        "a necktie accessory for clothing",
    ],
    "skateboard": [
        "a board with wheels for riding, not a surfboard",
        "a skateboard on ground, distinct from a snowboard",
        "a small wheeled board for skating",
    ],
    "surfboard": [
        "a long board for riding waves, not a skateboard",
        "a surfboard on water, distinct from a snowboard",
        "a surfing board, longer than a skateboard",
    ],
    "snowboard": [
        "a board for riding on snow, not a skateboard",
        "a snowboard on snow, distinct from skis",
        "a single wide board for snow sports",
    ],
    "skis": [
        "a pair of long narrow boards for snow, not a snowboard",
        "skis for skiing, two separate boards",
        "ski equipment distinct from a snowboard",
    ],
    "sports ball": [
        "a round ball for sports, not a frisbee",
        "a sports ball like soccer, basketball, or tennis ball",
        "a spherical ball distinct from an orange or apple",
    ],
    "kite": [
        "a flying object on a string in the sky",
        "a kite in the air, not a bird or airplane",
        "a colorful kite distinct from an umbrella",
    ],
    "baseball bat": [
        "a long stick for hitting baseballs",
        "a baseball bat, not a tennis racket",
        "a wooden or metal bat for baseball",
    ],
    "baseball glove": [
        "a leather glove for catching baseballs",
        "a baseball mitt, not a regular glove",
        "a catching glove for baseball",
    ],
    "tennis racket": [
        "a racket with strings for hitting tennis balls",
        "a tennis racket, not a baseball bat",
        "a racket with oval head and handle",
    ],
    "frisbee": [
        "a flat disc for throwing, not a sports ball",
        "a frisbee disc, not a plate or ball",
        "a flying disc toy distinct from a ball",
    ],
    "toaster": [
        "a small appliance for toasting bread",
        "a toaster, not a microwave or oven",
        "a countertop toaster distinct from an oven",
    ],
    "sink": [
        "a basin with faucet for washing",
        "a sink in kitchen or bathroom, not a toilet or bathtub",
        "a wash basin distinct from a bowl or toilet",
    ],
    "scissors": [
        "a cutting tool with two blades",
        "scissors for cutting, not a knife",
        "a pair of scissors with handles",
    ],
    "teddy bear": [
        "a stuffed plush toy bear",
        "a teddy bear toy, not a real bear",
        "a soft stuffed animal distinct from a real bear",
    ],
    "hair drier": [
        "an electric device for drying hair",
        "a hair dryer appliance, not a toaster",
        "a blow dryer for hair styling",
    ],
    "toothbrush": [
        "a small brush for cleaning teeth",
        "a toothbrush with bristles and handle",
        "a dental hygiene brush, not a hair brush",
    ],
    "vase": [
        "a decorative container for flowers",
        "a vase, not a bottle or cup",
        "a flower vase distinct from a bowl or bottle",
    ],
    "clock": [
        "a device showing time with hands or digits",
        "a clock on wall or table, not a watch",
        "a time-telling clock distinct from a TV screen",
    ],
    "traffic light": [
        "a signal light at intersections with red yellow green",
        "a traffic light, not a lamp or street light",
        "a traffic signal distinct from a regular light",
    ],
    "fire hydrant": [
        "a red water outlet on the street for firefighting",
        "a fire hydrant, not a parking meter",
        "a street hydrant distinct from a bollard",
    ],
    "stop sign": [
        "an octagonal red sign saying stop",
        "a stop sign, not a traffic light",
        "a road stop sign distinct from other signs",
    ],
    "parking meter": [
        "a coin-operated meter for parking time",
        "a parking meter, not a fire hydrant",
        "a street parking meter distinct from a hydrant",
    ],
    "elephant": [
        "a large gray animal with trunk and tusks",
        "an elephant, not a bear or horse, very large",
        "an elephant with long trunk distinct from other animals",
    ],
    "zebra": [
        "a striped black and white horse-like animal",
        "a zebra with stripes, not a horse",
        "a zebra distinct from a horse by its stripes",
    ],
    "giraffe": [
        "a tall animal with very long neck and spots",
        "a giraffe, not a horse or zebra, extremely tall",
        "a spotted giraffe with long neck",
    ],
    "bear": [
        "a large furry wild animal, not a teddy bear",
        "a real bear, not a stuffed toy bear",
        "a bear animal distinct from a dog or teddy bear",
    ],
    # ============ COCO-Stuff specific entries (no-space class names) ============
    "trafficlight": [
        "a signal light at intersections with red yellow green",
        "a traffic light, not a lamp or street light",
        "a traffic signal distinct from a regular light",
    ],
    "firehydrant": [
        "a red water outlet on the street for firefighting",
        "a fire hydrant, not a parking meter",
        "a street hydrant distinct from a bollard",
    ],
    "stopsign": [
        "an octagonal red sign saying stop",
        "a stop sign, not a traffic light",
        "a road stop sign distinct from a banner",
    ],
    "parkingmeter": [
        "a coin-operated meter for parking time",
        "a parking meter, not a fire hydrant",
        "a street parking meter on sidewalk",
    ],
    "sportsball": [
        "a round ball for sports, not a frisbee or fruit",
        "a sports ball like soccer, basketball, or tennis ball",
        "a spherical ball distinct from an orange or apple",
    ],
    "baseballbat": [
        "a long stick for hitting baseballs",
        "a baseball bat, not a tennis racket or branch",
        "a wooden or metal bat for baseball",
    ],
    "baseballglove": [
        "a leather glove for catching baseballs",
        "a baseball mitt, not a regular glove",
        "a catching glove for baseball",
    ],
    "tennisracket": [
        "a racket with strings for hitting tennis balls",
        "a tennis racket, not a baseball bat",
        "a racket with oval head and handle",
    ],
    "wineglass": [
        "a stemmed glass for drinking wine",
        "a wine glass with stem, not a regular cup or bowl",
        "a tall stemmed glass distinct from a cup",
    ],
    "hotdog": [
        "a sausage in a bun, not a sandwich or pizza",
        "a hot dog with bread roll, distinct from a sandwich",
        "a hot dog food item with elongated shape",
    ],
    "diningtable": [
        "a table used for eating meals, not a desk or counter",
        "a dining table with flat surface, distinct from a desk",
        "a table for dining, often with chairs around it",
    ],
    "cellphone": [
        "a handheld mobile phone or smartphone",
        "a cell phone with screen, not a remote control",
        "a smartphone device distinct from a remote",
    ],
    "teddybear": [
        "a stuffed plush toy bear",
        "a teddy bear toy, not a real bear",
        "a soft stuffed animal distinct from a real bear",
    ],
    "hairdrier": [
        "an electric device for drying hair",
        "a hair dryer appliance, not a toaster",
        "a blow dryer for hair styling",
    ],
    "toilet": [
        "a bathroom fixture for waste disposal",
        "a toilet bowl, not a sink or bathtub",
        "a white porcelain toilet in bathroom",
    ],
    # ============ COCO-Stuff confusable stuff categories ============
    "building-other": [
        "a building structure, not a house or skyscraper",
        "a generic building distinct from a house",
        "a building that is not residential",
    ],
    "ceiling-other": [
        "an indoor ceiling surface, not tiled",
        "a plain ceiling, distinct from ceiling tiles",
        "a ceiling surface in a room",
    ],
    "ceiling-tile": [
        "a tiled ceiling with grid pattern",
        "ceiling tiles in a drop ceiling, not a plain ceiling",
        "a ceiling with tile panels",
    ],
    "wall-brick": [
        "a wall made of bricks",
        "a brick wall with visible mortar joints",
        "a brick wall distinct from concrete or wood walls",
    ],
    "wall-concrete": [
        "a wall made of concrete",
        "a smooth concrete wall, not brick or stone",
        "a concrete wall surface",
    ],
    "wall-other": [
        "a wall of unspecified material",
        "a generic wall, not brick, concrete, stone, tile, or wood",
        "a wall surface of other material",
    ],
    "wall-stone": [
        "a wall made of stone blocks",
        "a stone wall, not brick or concrete",
        "a wall with natural stone surface",
    ],
    "wall-tile": [
        "a wall covered with tiles",
        "a tiled wall in bathroom or kitchen",
        "a wall with ceramic or porcelain tiles",
    ],
    "wall-wood": [
        "a wall made of wood panels",
        "a wooden wall, not brick or concrete",
        "a wall with wood paneling",
    ],
    "wall-panel": [
        "a wall with panel sections",
        "a paneled wall, distinct from plain wall",
        "a wall with decorative panels",
    ],
    "floor-marble": [
        "a floor made of marble",
        "a marble floor with polished surface",
        "a marble floor distinct from tile or wood",
    ],
    "floor-other": [
        "a floor of unspecified material",
        "a generic floor surface",
        "a floor that is not marble, stone, tile, or wood",
    ],
    "floor-stone": [
        "a floor made of stone",
        "a stone floor, not marble or tile",
        "a floor with natural stone surface",
    ],
    "floor-tile": [
        "a floor covered with tiles",
        "a tiled floor, not marble or wood",
        "a floor with ceramic tiles",
    ],
    "floor-wood": [
        "a floor made of wood",
        "a wooden floor with planks or parquet",
        "a hardwood floor distinct from tile",
    ],
    "sky-other": [
        "sky that is not clear blue, overcast or twilight",
        "a sky area, distinct from clouds",
        "the sky background in a scene",
    ],
    "window-blind": [
        "window blinds or shades",
        "blinds covering a window, not curtains",
        "window blinds distinct from curtains",
    ],
    "window-other": [
        "a window that is not covered by blinds",
        "a glass window, not a window blind",
        "a window opening in a building",
    ],
    "door-stuff": [
        "a door surface or structure",
        "a door, not a window or wall",
        "a door in a building",
    ],
    "desk-stuff": [
        "a desk surface or structure",
        "a desk, not a dining table or counter",
        "a work desk distinct from a table",
    ],
    "mirror-stuff": [
        "a reflective mirror surface",
        "a mirror, not a window or screen",
        "a mirror reflecting the scene",
    ],
    "curtain": [
        "fabric hanging over a window",
        "curtains or drapes, not window blinds",
        "window curtains distinct from cloth or towel",
    ],
    "carpet": [
        "a floor covering made of fabric",
        "a carpet, not a rug or mat",
        "a wall-to-wall carpet on the floor",
    ],
    "clothes": [
        "clothing items or garments",
        "clothes worn by people or hanging",
        "clothing distinct from cloth or textile",
    ],
    "clouds": [
        "white or gray clouds in the sky",
        "clouds floating in the sky, not fog",
        "cloud formations distinct from sky",
    ],
    "dirt": [
        "bare dirt or soil on the ground",
        "dirt ground, not mud or gravel",
        "a dirt surface distinct from pavement",
    ],
    "gravel": [
        "small loose stones on the ground",
        "gravel surface, not sand or dirt",
        "a gravel path or area",
    ],
    "sand": [
        "sandy surface on beach or desert",
        "sand on the ground, not dirt or gravel",
        "a sandy terrain distinct from soil",
    ],
    "pavement": [
        "a paved surface for walking or driving",
        "pavement or sidewalk, not road",
        "a paved area distinct from road surface",
    ],
    "playingfield": [
        "a sports field or playing area",
        "a playing field with grass or turf",
        "a sports field distinct from regular grass",
    ],
    "railroad": [
        "railway tracks and rail infrastructure",
        "a railroad with tracks, not a road",
        "railroad tracks distinct from a road",
    ],
    "house": [
        "a residential building or home",
        "a house, not a generic building or skyscraper",
        "a residential house distinct from other buildings",
    ],
    "roof": [
        "the top covering of a building",
        "a roof surface, not a ceiling",
        "a building roof seen from outside",
    ],
    "stairs": [
        "steps for going up or down",
        "stairs or staircase, not a ramp",
        "a flight of stairs in a building",
    ],
    "stone": [
        "natural stone material or rocks",
        "stone surface, not concrete or brick",
        "natural stone distinct from gravel",
    ],
    "fruit": [
        "edible fruit, not a specific named fruit",
        "generic fruit, not banana, apple, or orange specifically",
        "fruit items distinct from vegetables",
    ],
    "vegetable": [
        "edible vegetable, not a specific named vegetable",
        "generic vegetable, not broccoli or carrot specifically",
        "vegetable items distinct from fruit",
    ],
    "food-other": [
        "food that is not a specific named food item",
        "generic food, not pizza, cake, or sandwich",
        "other food items on a plate or table",
    ],
    "furniture-other": [
        "furniture that is not a chair, couch, bed, or table",
        "generic furniture piece",
        "other furniture distinct from named items",
    ],
    "ground-other": [
        "ground surface not otherwise classified",
        "generic ground, not grass, dirt, or pavement",
        "other ground surface",
    ],
    "solid-other": [
        "solid material not otherwise classified",
        "generic solid surface or object",
        "other solid material",
    ],
    "structural-other": [
        "structural element not otherwise classified",
        "generic structural component of a building",
        "other structural element",
    ],
    "textile-other": [
        "textile material not otherwise classified",
        "generic fabric or textile",
        "other textile distinct from cloth or clothes",
    ],
    "water-other": [
        "water that is not sea, river, or lake",
        "generic water surface or body",
        "other water distinct from named water bodies",
    ],
    "plant-other": [
        "a plant not otherwise classified",
        "generic plant, not a potted plant or tree",
        "other plant vegetation",
    ],
    "cloth": [
        "a piece of fabric or cloth material",
        "cloth material, not clothes or curtain",
        "a cloth surface distinct from carpet or towel",
    ],
    "leaves": [
        "tree or plant leaves",
        "leaves on branches or ground, not a whole tree",
        "foliage leaves distinct from grass",
    ],
    "bush": [
        "a shrub or bush plant",
        "a bush, not a tree or potted plant",
        "a low shrub distinct from a tree",
    ],
    "moss": [
        "green moss growing on surfaces",
        "moss on rocks or ground, not grass",
        "moss vegetation distinct from grass",
    ],
    "fog": [
        "atmospheric fog or mist",
        "fog in the air, not clouds",
        "foggy haze distinct from clouds in sky",
    ],
    "mud": [
        "wet soil or mud on the ground",
        "mud, not dirt or water",
        "muddy ground distinct from dry dirt",
    ],
    "snow": [
        "white snow on the ground or surfaces",
        "snow covering, not ice or clouds",
        "snow on ground distinct from sky",
    ],
    "railing": [
        "a handrail or guardrail barrier",
        "a railing, not a fence",
        "a railing along stairs or edges",
    ],
    "bridge": [
        "a structure spanning over water or road",
        "a bridge, not a road or platform",
        "a bridge structure for crossing",
    ],
    "banner": [
        "a hanging banner or sign",
        "a banner, not a curtain or flag",
        "a decorative or informational banner",
    ],
    "tent": [
        "a portable shelter made of fabric",
        "a tent, not a building or house",
        "a camping or event tent",
    ],
    "towel": [
        "a cloth for drying, not a napkin or blanket",
        "a towel in bathroom or kitchen",
        "a towel distinct from cloth or rug",
    ],
    "napkin": [
        "a small cloth or paper for wiping",
        "a napkin, not a towel or cloth",
        "a table napkin distinct from a towel",
    ],
    "cardboard": [
        "cardboard material or boxes",
        "cardboard, not paper or wood",
        "a cardboard surface or box",
    ],
    "net": [
        "a mesh net structure",
        "a net, not a fence or cage",
        "a net for sports or fishing",
    ],
    "cage": [
        "a barred enclosure or cage",
        "a cage, not a fence or net",
        "a cage structure with bars",
    ],
    "counter": [
        "a flat surface for working or serving",
        "a counter in kitchen or store, not a table",
        "a countertop surface distinct from a desk",
    ],
    "cupboard": [
        "a storage cabinet with doors",
        "a cupboard, not a shelf or wardrobe",
        "a kitchen or room cupboard",
    ],
    "shelf": [
        "a flat surface for storing items",
        "a shelf on a wall, not a counter or table",
        "shelves for storage or display",
    ],
    "mat": [
        "a small floor covering",
        "a mat, not a carpet or rug",
        "a floor mat distinct from a carpet",
    ],
    "metal": [
        "metallic material or surface",
        "metal surface, not plastic or wood",
        "a metal material or object",
    ],
    "plastic": [
        "plastic material or surface",
        "plastic, not metal or glass",
        "a plastic material or object",
    ],
    "paper": [
        "paper material or sheets",
        "paper, not cardboard or cloth",
        "paper surface or documents",
    ],
    "straw": [
        "dried plant stalks or straw material",
        "straw, not grass or hay",
        "straw material on ground or in bales",
    ],
}

def get_contrastive_prompts(class_name: str):
    """Return contrastive descriptions that improve class discrimination."""
    c = class_name.replace("_", " ").strip().lower()
    c_nospace = class_name.replace("_", "").strip().lower()
    return SEMANTIC_RELATIONS.get(c, None) or SEMANTIC_RELATIONS.get(c_nospace, [])

@MODELS.register_module()
class NACLIP(BaseSegmentor):
    def __init__(self, clip_path, name_path, device=torch.device('cuda'),
                 arch='reduced', attn_strategy='naclip', gaussian_std=5.,
                 adaptive_neighbour=False, sigma_min=1.5, sigma_max=7.0, sigma_alpha=4.0, neigh_k=4,
                 omega_scale_mode='none', omega_scale=None, kk_norm=False,
                 use_domain_prompts=False,
                 tfd_beta=0.0,
                 attn_temp=1.0,
                 ssca_temp=0.0, ssca_lambda=1.0,
                 cte_alpha=0.0,  # CTE: Contrastive Text Enhancement
                 mlfa_layers=1, mlfa_weights=None,
                 residual_alpha=0.0, ffn_alpha=0.0,
                 sharp_gamma=0.0,
                 lla_tau=0.0,
                 lla_mode='logsumexp',
                 pamr_steps=10, pamr_stride=(8, 16),
                 prob_thd=0.0, logit_scale=40, slide_stride=112, slide_crop=224):

        # Handle omega_scale as alias for omega_scale_mode (config compatibility)
        if omega_scale is not None and omega_scale_mode == 'none':
            omega_scale_mode = omega_scale

        data_preprocessor = SegDataPreProcessor(mean=[122.771, 116.746, 104.094], std=[68.501, 66.632, 70.323], rgb_to_bgr=True)
        super().__init__(data_preprocessor=data_preprocessor)
        self.net, _ = clip.load(clip_path, device=device, jit=False)

        query_words, self.query_idx = get_cls_idx(name_path)
        self.num_queries = len(query_words)
        self.query_words = [qw.strip() for qw in query_words]
        self.query_idx = torch.Tensor(self.query_idx).to(torch.int64).to(device)


        self.use_domain_prompts = use_domain_prompts
        self.cte_alpha = float(cte_alpha)  # CTE blending weight

        # Detect dataset for domain prompt routing
        self.dataset_name = detect_dataset_from_path(name_path)
        if self.use_domain_prompts:
            logging.info(f'[NACLIP] Domain prompts enabled for dataset: {self.dataset_name}')

        query_features = []
        contrastive_features = []  # For CTE

        with torch.no_grad():
            for qw in query_words:
                # Standard text features
                if self.use_domain_prompts:
                    prompts = get_prompt_bank(qw, self.dataset_name)
                else:
                    c = qw.replace("_", " ").strip().lower()
                    prompts = [temp(c) for temp in openai_imagenet_template]
                query = clip.tokenize(prompts).to(device)
                feats = self.net.encode_text(query)
                feats = feats / feats.norm(dim=-1, keepdim=True)
                feature = feats.mean(dim=0)
                feature = feature / feature.norm()
                query_features.append(feature.unsqueeze(0))

                # CTE: Contrastive Text Enhancement
                # Encode contrastive descriptions for confusable classes.
                contrastive_prompts = get_contrastive_prompts(qw)
                if contrastive_prompts and self.cte_alpha > 0:
                    cte_query = clip.tokenize(contrastive_prompts).to(device)
                    cte_feats = self.net.encode_text(cte_query)
                    cte_feats = cte_feats / cte_feats.norm(dim=-1, keepdim=True)
                    cte_feature = cte_feats.mean(dim=0)
                    cte_feature = cte_feature / cte_feature.norm()
                    contrastive_features.append(cte_feature.unsqueeze(0))
                else:
                    # Fall back to the original feature when no description exists.
                    contrastive_features.append(feature.unsqueeze(0))

        self.query_features = torch.cat(query_features, dim=0)

        # Blend the original and contrastive text features.
        if self.cte_alpha > 0:
            contrastive_features = torch.cat(contrastive_features, dim=0)
            # Weighted blend: (1-alpha) * original + alpha * contrastive.
            self.query_features = (1 - self.cte_alpha) * self.query_features + self.cte_alpha * contrastive_features
            self.query_features = self.query_features / self.query_features.norm(dim=-1, keepdim=True)
            logging.info(f'[NACLIP] CTE applied: alpha={self.cte_alpha}')

        # TFD: Text Feature Decorrelation
        if tfd_beta > 0:
            T = self.query_features  # (C, D), already L2-normalized
            # Gram matrix of pairwise cosine similarities
            G = T @ T.T  # (C, C)
            # Zero out self-similarity so we only subtract cross-class components
            G = G - torch.diag(G.diag())
            # t_i' = t_i - beta * sum_j (t_i . t_j) * t_j
            T_decorr = T - tfd_beta * (G @ T)
            # Re-normalize
            self.query_features = T_decorr / T_decorr.norm(dim=-1, keepdim=True)
            logging.info(f'[NACLIP] TFD applied: beta={tfd_beta}')

        self.dtype = self.query_features.dtype
        if attn_strategy == 'clip' and arch != 'vanilla':
            logging.warning(f'attn_strategy=clip implies arch=vanilla, but got arch={arch}. Auto set arch=vanilla.')
            arch = 'vanilla'

        self.net.visual.set_params(arch, attn_strategy, gaussian_std,
                                  adaptive_neighbour=adaptive_neighbour, sigma_min=sigma_min, sigma_max=sigma_max, sigma_alpha=sigma_alpha, neigh_k=neigh_k,
                                  omega_scale_mode=omega_scale_mode, kk_norm=kk_norm,
                                  attn_temp=attn_temp,
                                  mlfa_layers=mlfa_layers, mlfa_weights=mlfa_weights,
                                  residual_alpha=residual_alpha, ffn_alpha=ffn_alpha)
        self.logit_scale = logit_scale
        self.prob_thd = prob_thd
        self.slide_stride = slide_stride
        self.slide_crop = slide_crop
        self.align_corners = False
        self.pamr = PAMR(pamr_steps, dilations=pamr_stride).to(device) if pamr_steps > 0 else None

        # SSCA: Spatial-Semantic Cross-Attention
        self.ssca_temp = float(ssca_temp)
        self.ssca_lambda = float(ssca_lambda)
        self.tfd_beta = tfd_beta  # Store for potential use in SSCA
        self.sharp_gamma = float(sharp_gamma)  # Logit spatial sharpening
        self.lla_tau = float(lla_tau)  # LLA: Logit-Level Aggregation temperature
        self.lla_mode = str(lla_mode).lower() if lla_mode is not None else 'logsumexp'
        if self.lla_mode not in {'off', 'logsumexp', 'max', 'mean'}:
            raise ValueError(f'Unsupported lla_mode: {lla_mode}')

        logging.info(f'[NACLIP] arch={arch}, attn_strategy={attn_strategy}, gaussian_std={gaussian_std}, ' \
                f'adaptive_neighbour={adaptive_neighbour}, sigma_min={sigma_min}, sigma_max={sigma_max}, sigma_alpha={sigma_alpha}, neigh_k={neigh_k}, ' \
                f'tfd_beta={tfd_beta}, attn_temp={attn_temp}, ssca_temp={ssca_temp}, ssca_lambda={ssca_lambda}, cte_alpha={cte_alpha}, ' \
                f'mlfa_layers={mlfa_layers}, mlfa_weights={mlfa_weights}, ' \
                f'lla_mode={self.lla_mode}, lla_tau={self.lla_tau}, '
                f'pamr_steps={pamr_steps}, slide_stride={slide_stride}, slide_crop={slide_crop}')



    def forward_feature(self, img):
            if type(img) == list:
                img = img[0]

            tokens = self.net.encode_image(img, return_all=True)  # (B, 1+N, D)
            patch_feats = tokens[:, 1:]

            patch_feats = patch_feats / patch_feats.norm(dim=-1, keepdim=True)

            # Base similarity
            base_logits = patch_feats @ self.query_features.T  # (B, N, C)

            # SSCA: Spatial-Semantic Cross-Attention
            if self.ssca_temp > 0:
                logits = self._apply_ssca(patch_feats, base_logits)
            else:
                logits = base_logits

            patch_size = self.net.visual.patch_size
            w, h = img[0].shape[-2] // patch_size, img[0].shape[-1] // patch_size
            out_dim = logits.shape[-1]
            logits = logits.permute(0, 2, 1).reshape(-1, out_dim, w, h)

            # Spatial sharpening: unsharp mask on logits
            if self.sharp_gamma > 0:
                # 3x3 average pooling as blur, then sharpen
                blurred = F.avg_pool2d(logits, kernel_size=3, stride=1, padding=1)
                logits = logits + self.sharp_gamma * (logits - blurred)

            logits = nn.functional.interpolate(logits, size=img.shape[-2:], mode='bilinear', align_corners=self.align_corners)
            return logits

    def _apply_ssca(self, patch_feats, base_logits):
        """
        Spatial-Semantic Cross-Attention (SSCA)

        Convert static text features into image-adaptive contextualized features.

        The base logits identify relevant patches for each class. Their visual
        features are pooled into image-specific prototypes, which produce
        refined logits that are blended with the base prediction.
        """
        # Step 1: Text-to-Patch Cross-Attention
        attn_weights = F.softmax(base_logits.permute(0, 2, 1) / self.ssca_temp, dim=-1)  # (B, C, N)

        # Step 2: Generate Contextualized Text Features
        ctx_text = torch.bmm(attn_weights, patch_feats)  # (B, C, D)
        ctx_text = F.normalize(ctx_text, dim=-1)

        # Step 3: Compute Refined Similarity
        refined_logits = torch.bmm(patch_feats, ctx_text.permute(0, 2, 1))  # (B, N, C)

        # Step 4: Blend base and refined logits
        logits = base_logits + self.ssca_lambda * (refined_logits - base_logits)

        return logits

    def forward_slide(self, img, stride=112, crop_size=224):
        """
        Inference by sliding-window with overlap. If h_crop > h_img or w_crop > w_img,
        the small patch will be used to decode without padding.
        """
        if type(img) == list:
            img = img[0].unsqueeze(0)
        if type(stride) == int:
            stride = (stride, stride)
        if type(crop_size) == int:
            crop_size = (crop_size, crop_size)

        h_stride, w_stride = stride
        h_crop, w_crop = crop_size
        batch_size, _, h_img, w_img = img.shape
        out_channels = self.num_queries
        h_grids = max(h_img - h_crop + h_stride - 1, 0) // h_stride + 1
        w_grids = max(w_img - w_crop + w_stride - 1, 0) // w_stride + 1
        preds = img.new_zeros((batch_size, out_channels, h_img, w_img))
        count_mat = img.new_zeros((batch_size, 1, h_img, w_img))
        for h_idx in range(h_grids):
            for w_idx in range(w_grids):
                y1 = h_idx * h_stride
                x1 = w_idx * w_stride
                y2 = min(y1 + h_crop, h_img)
                x2 = min(x1 + w_crop, w_img)
                y1 = max(y2 - h_crop, 0)
                x1 = max(x2 - w_crop, 0)
                crop_img = img[:, :, y1:y2, x1:x2]
                crop_seg_logit = self.forward_feature(crop_img)
                preds += nn.functional.pad(crop_seg_logit,
                                           (int(x1), int(preds.shape[3] - x2), int(y1), int(preds.shape[2] - y2)))

                count_mat[:, :, y1:y2, x1:x2] += 1
        assert (count_mat == 0).sum() == 0

        logits = preds / count_mat
        return logits

    def predict(self, inputs, data_samples):
        if data_samples is not None:
            batch_img_metas = [data_sample.metainfo for data_sample in data_samples]
        else:
            batch_img_metas = [dict(
                ori_shape=inputs.shape[2:],
                img_shape=inputs.shape[2:],
                pad_shape=inputs.shape[2:],
                padding_size=[0, 0, 0, 0])
            ] * inputs.shape[0]

        if self.slide_crop > 0:
            seg_logits = self.forward_slide(inputs, self.slide_stride, self.slide_crop)
        else:
            seg_logits = self.forward_feature(inputs)

        img_size = batch_img_metas[0]['ori_shape']
        seg_logits = nn.functional.interpolate(seg_logits, size=img_size, mode='bilinear', align_corners=self.align_corners)

        if self.pamr:
            img = nn.functional.interpolate(inputs, size=img_size, mode='bilinear', align_corners=self.align_corners)
            try:
                seg_logits = self.pamr(img, seg_logits.to(img.dtype)).to(self.dtype)
            except RuntimeError as e:
                logging.warning(f"Couldn't apply PAMR for image {batch_img_metas[0]['img_path'].split('/')[-1]} "
                                f"of size {img_size}, probably due to low memory. Error message: \"{str(e)}\"")

        return self.postprocess_result(seg_logits, data_samples)

    def postprocess_result(self, seg_logits, data_samples):
        batch_size = seg_logits.shape[0]
        for i in range(batch_size):
            num_cls, num_queries = max(self.query_idx) + 1, len(self.query_idx)

            use_lla = (
                num_cls != num_queries and
                self.lla_mode != 'off' and
                (self.lla_mode in {'max', 'mean'} or self.lla_tau > 0)
            )

            if use_lla:
                # LLA: Logit-Level Aggregation (Background-Selective)
                # Aggregate background sub-embeddings BEFORE softmax;
                # foreground multi-embedding classes keep max aggregation.
                raw_logits = seg_logits[i]  # (num_queries, H, W)

                # Check if background class (idx 0) has multiple sub-embeddings
                bg_mask = (self.query_idx == 0)
                bg_count = bg_mask.sum().item()

                if bg_count > 1:
                    # Step 1: Aggregate background sub-embeddings via the selected reducer.
                    bg_logits = raw_logits[bg_mask]  # (K_bg, H, W)
                    if self.lla_mode == 'logsumexp':
                        bg_agg = self.lla_tau * torch.logsumexp(bg_logits / self.lla_tau, dim=0, keepdim=True)
                    elif self.lla_mode == 'max':
                        bg_agg = bg_logits.max(dim=0, keepdim=True)[0]
                    elif self.lla_mode == 'mean':
                        bg_agg = bg_logits.mean(dim=0, keepdim=True)
                    else:
                        raise RuntimeError(f'Invalid lla_mode at runtime: {self.lla_mode}')

                    # Step 2: Build new logits tensor: bg_agg + all non-bg queries
                    fg_mask = ~bg_mask
                    fg_logits = raw_logits[fg_mask]  # (num_fg_queries, H, W)
                    new_logits = torch.cat([bg_agg, fg_logits], dim=0)  # (1 + num_fg_queries, H, W)

                    # Step 3: softmax over new logits
                    seg_probs_all = torch.softmax(new_logits * self.logit_scale, dim=0)

                    # Step 4: max-aggregate fg classes (memory-efficient, per-class)
                    fg_idx = self.query_idx[fg_mask]
                    seg_probs = seg_probs_all.new_zeros(num_cls, *seg_probs_all.shape[1:])
                    seg_probs[0] = seg_probs_all[0]  # background (already single)
                    for c in range(1, num_cls):
                        c_mask = (fg_idx == c)
                        if c_mask.any():
                            # +1 offset because seg_probs_all[0] is bg
                            c_probs = seg_probs_all[1:][c_mask]  # (K_c, H, W)
                            seg_probs[c] = c_probs.max(0)[0]
                else:
                    # No multi-embedding background, fall through to standard path
                    seg_probs = torch.softmax(seg_logits[i] * self.logit_scale, dim=0)
                    if num_cls != num_queries:
                        seg_probs = seg_probs.unsqueeze(0)
                        cls_index = nn.functional.one_hot(self.query_idx)
                        cls_index = cls_index.T.view(num_cls, num_queries, 1, 1)
                        seg_probs = (seg_probs * cls_index).max(1)[0]
            else:
                seg_probs = torch.softmax(seg_logits[i] * self.logit_scale, dim=0)  # n_queries * w * h

                if num_cls != num_queries:
                    seg_probs = seg_probs.unsqueeze(0)
                    cls_index = nn.functional.one_hot(self.query_idx)
                    cls_index = cls_index.T.view(num_cls, num_queries, 1, 1)
                    seg_probs = (seg_probs * cls_index).max(1)[0]

            seg_pred = seg_probs.argmax(0, keepdim=True)
            seg_pred[seg_probs.max(0, keepdim=True)[0] < self.prob_thd] = 0
            seg_probs /= seg_probs.sum(0, keepdim=True)

            data_samples[i].set_data({
                'seg_logits': PixelData(**{'data': seg_probs}),
                'pred_sem_seg': PixelData(**{'data': seg_pred})
            })

        return data_samples

    def _forward(data_samples):
        pass

    def inference(self, img, batch_img_metas):
        pass

    def encode_decode(self, inputs, batch_img_metas):
        pass

    def extract_feat(self, inputs):
        pass

    def loss(self, inputs, data_samples):
        pass


def get_cls_idx(path):
    with open(path, 'r') as f:
        name_sets = f.readlines()
    num_cls = len(name_sets)

    class_names, class_indices = list(), list()
    for idx in range(num_cls):
        names_i = name_sets[idx].split(', ')
        class_names += names_i
        class_indices += [idx for _ in range(len(names_i))]
    class_names = [item.replace('\n', '') for item in class_names]
    return class_names, class_indices