mirror of
https://github.com/invoke-ai/InvokeAI
synced 2024-08-30 20:32:17 +00:00
Simplifying ControlNet SAM preprocessor segmentation color mapping.
This commit is contained in:
parent
de4064bdac
commit
b872e7a5e0
@ -485,18 +485,12 @@ class SamDetectorReproducibleColors(SamDetector):
|
||||
sorted_anns = sorted(anns, key=(lambda x: x['area']), reverse=True)
|
||||
h, w = anns[0]['segmentation'].shape
|
||||
final_img = Image.fromarray(np.zeros((h, w, 3), dtype=np.uint8), mode="RGB")
|
||||
print("number of annotations: ", len(sorted_anns))
|
||||
print("type of annotations: ", type(sorted_anns))
|
||||
palette = ade_palette()
|
||||
for i, ann in enumerate(sorted_anns):
|
||||
m = ann['segmentation']
|
||||
img = np.empty((m.shape[0], m.shape[1], 3), dtype=np.uint8)
|
||||
# doing modulo just in case number of annotated regions exceeds number of colors in palette
|
||||
ann_color = palette[i % len(palette)]
|
||||
print(ann_color)
|
||||
img[:, :, 0] = ann_color[0]
|
||||
img[:, :, 1] = ann_color[1]
|
||||
img[:, :, 2] = ann_color[2]
|
||||
img[:, :] = ann_color
|
||||
final_img.paste(Image.fromarray(img, mode="RGB"), (0, 0), Image.fromarray(np.uint8(m * 255)))
|
||||
|
||||
return np.array(final_img, dtype=np.uint8)
|
||||
|
Loading…
Reference in New Issue
Block a user