fix labels

app.py

@@ -4,7 +4,7 @@ import torch
 from torchvision.transforms import InterpolationMode
 
 BICUBIC = InterpolationMode.BICUBIC
-from utils import setup, get_similarity_map, display_segmented_sketch,get_noun_phrase
+from utils import setup, get_similarity_map,get_noun_phrase, rgb_to_hsv, hsv_to_rgb
 from vpt.launch import default_argument_parser
 from collections import OrderedDict
 import numpy as np
@@ -81,8 +81,52 @@ def run(sketch, caption, threshold, seed):
     pixel_similarity[pixel_similarity < threshold] = 0
     pixel_similarity_array = pixel_similarity.cpu().numpy().transpose(2, 0, 1)
 
-    display_segmented_sketch(pixel_similarity_array, sketch2, classes, classes_colors, live=True)
+    # display_segmented_sketch(pixel_similarity_array, sketch2, classes, classes_colors, live=True)
+    
+    # Find the class index with the highest similarity for each pixel
+    class_indices = np.argmax(pixel_similarity_array, axis=0)
+    # Create an HSV image placeholder
+    hsv_image = np.zeros(class_indices.shape + (3,))  # Shape (512, 512, 3)
+    hsv_image[..., 2] = 1  # Set Value to 1 for a white base
+    
+    # Set the hue and value channels
+    for i, color in enumerate(classes_colors):
+        rgb_color = np.array(color).reshape(1, 1, 3)
+        hsv_color = rgb_to_hsv(rgb_color)
+        mask = class_indices == i
+        if i < len(classes):  # For the first N-2 classes, set color based on similarity
+            hsv_image[..., 0][mask] = hsv_color[0, 0, 0]  # Hue
+            hsv_image[..., 1][mask] = pixel_similarity_array[i][mask] > 0  # Saturation
+            hsv_image[..., 2][mask] = pixel_similarity_array[i][mask]  # Value
+        else:  # For the last two classes, set pixels to black
+            hsv_image[..., 0][mask] = 0  # Hue doesn't matter for black
+            hsv_image[..., 1][mask] = 0  # Saturation set to 0
+            hsv_image[..., 2][mask] = 0  # Value set to 0, making it black
+    
+    mask_tensor_org = sketch2[:,:,0]/255
+    hsv_image[mask_tensor_org==1] = [0,0,1]
 
+    # Convert the HSV image back to RGB to display and save
+    rgb_image = hsv_to_rgb(hsv_image)
+
+    
+    if len(classes) > 1:
+        # Calculate centroids and render class names
+        for i, class_name in enumerate(classes):
+            mask = class_indices == i
+            if np.any(mask):
+                y, x = np.nonzero(mask)
+                centroid_x, centroid_y = np.mean(x), np.mean(y)
+                plt.text(centroid_x, centroid_y, class_name, color=classes_colors[i], ha='center', va='center',fontsize=10,   # color=classes_colors[i]
+                bbox=dict(facecolor='lightgrey', edgecolor='none', boxstyle='round,pad=0.2', alpha=0.8))
+
+    # Display the image with class names
+    plt.imshow(rgb_image)
+    plt.axis('off')
+    plt.tight_layout()
+    plt.savefig('output.png', bbox_inches='tight', pad_inches=0)
+    plt.close()
+    
     rgb_image = Image.open('output.png')
 
     return rgb_image