BrandSmith / src /models.py
raznis's picture
Upload folder using huggingface_hub
c9f5fa7 verified
raw
history blame contribute delete
677 Bytes
from pydantic import BaseModel
from PIL import Image
class ColorItem(BaseModel):
R: int
"The red component of the color."
G: int
"The green component of the color."
B: int
"The blue component of the color."
count: int
"The frequency of the color in the image."
class ColorPalette(BaseModel):
colors: list[ColorItem]
"""A list of colors and their frequency."""
def to_rgb(self):
return [(color.R, color.G, color.B) for color in self.colors]
def to_hex(self):
return [f"#{color.R:02x}{color.G:02x}{color.B:02x}" for color in self.colors]
def get_num_colors(self):
return len(self.colors)