DoruC/Grounded-Segment-Anything

1import torch2import torch.nn as nn3import torch.optim as optim4import numpy as np5import torch.nn.functional as F6 7 8class MLP(nn.Module):9    def __init__(self, input_size, hidden_size, num_classes, dropout_prob=0.1):10        super(MLP, self).__init__()11        self.fc1 = nn.Linear(input_size, hidden_size)12        self.relu = nn.ReLU()13        self.dropout = nn.Dropout(dropout_prob)14        self.fc2 = nn.Linear(hidden_size, num_classes)15 16    def forward(self, x):17        out = self.fc1(x)18        out = self.relu(out)19        out = self.dropout(out)20        out = self.fc2(out)21        return out22 23 24def show_anns(anns, color_code='auto'):25    if len(anns) == 0:26        return27    sorted_anns = sorted(anns, key=(lambda x: x['area']), reverse=True)28    ax = plt.gca()29    ax.set_autoscale_on(False)30    polygons = []31    color = []32    for ann in sorted_anns:33        m = ann['segmentation']34        img = np.ones((m.shape[0], m.shape[1], 3))35        color_mask = np.random.random((1, 3)).tolist()[0]36        if color_code == 'auto':37            for i in range(3):38                img[:,:,i] = color_mask[i]39        elif color_code == 'red':40            for i in range(3):41                img[:,:,0] = 142                img[:,:,1] = 043                img[:,:,2] = 044        else:45            for i in range(3):46                img[:,:,0] = 047                img[:,:,1] = 048                img[:,:,2] = 149    return np.dstack((img, m*0.35))50 51 52def show_points(coords, labels, ax, marker_size=375):53    pos_points = coords[labels==1]54    neg_points = coords[labels==0]55    ax.scatter(pos_points[:, 0], pos_points[:, 1], color='green', marker='*', 56               s=marker_size, edgecolor='white', linewidth=1.25)57    ax.scatter(neg_points[:, 0], neg_points[:, 1], color='red', marker='*', 58               s=marker_size, edgecolor='white', linewidth=1.25)   59 60def ram_show_mask(m):61    img = np.ones((m.shape[0], m.shape[1], 3))62    color_mask = np.random.random((1, 3)).tolist()[0]63    for i in range(3):64        img[:,:,0] = 165        img[:,:,1] = 066        img[:,:,2] = 067 68    return np.dstack((img, m*0.35))69 70 71def iou(mask1, mask2):72    intersection = np.logical_and(mask1, mask2)73    union = np.logical_or(mask1, mask2)74    iou_score = np.sum(intersection) / np.sum(union)75    return iou_score76 77 78def sort_and_deduplicate(sam_masks, iou_threshold=0.8):79    # Sort the sam_masks list based on the area value80    sorted_masks = sorted(sam_masks, key=lambda x: x['area'], reverse=True)81 82    # Deduplicate masks based on the given iou_threshold83    filtered_masks = []84    for mask in sorted_masks:85        duplicate = False86        for filtered_mask in filtered_masks:87            if iou(mask['segmentation'], filtered_mask['segmentation']) > iou_threshold:88                duplicate = True89                break90 91        if not duplicate:92            filtered_masks.append(mask)93 94    return filtered_masks95 96 97relation_classes = ['over',98        'in front of',99        'beside',100        'on',101        'in',102        'attached to',103        'hanging from',104        'on back of',105        'falling off',106        'going down',107        'painted on',108        'walking on',109        'running on',110        'crossing',111        'standing on',112        'lying on',113        'sitting on',114        'flying over',115        'jumping over',116        'jumping from',117        'wearing',118        'holding',119        'carrying',120        'looking at',121        'guiding',122        'kissing',123        'eating',124        'drinking',125        'feeding',126        'biting',127        'catching',128        'picking',129        'playing with',130        'chasing',131        'climbing',132        'cleaning',133        'playing',134        'touching',135        'pushing',136        'pulling',137        'opening',138        'cooking',139        'talking to',140        'throwing',141        'slicing',142        'driving',143        'riding',144        'parked on',145        'driving on',146        'about to hit',147        'kicking',148        'swinging',149        'entering',150        'exiting',151        'enclosing',152        'leaning on',]153