mirror of
https://github.com/AUTOMATIC1111/stable-diffusion-webui.git
synced 2024-12-21 07:30:02 +08:00
26 lines
703 B
Python
26 lines
703 B
Python
|
import os.path
|
||
|
|
import csv
|
||
|
|
from collections import namedtuple
|
||
|
|
|
||
|
|
Artist = namedtuple("Artist", ['name', 'weight', 'category'])
|
||
|
|
|
||
|
|
|
||
|
|
class ArtistsDatabase:
|
||
|
|
def __init__(self, filename):
|
||
|
|
self.cats = set()
|
||
|
|
self.artists = []
|
||
|
|
|
||
|
|
if not os.path.exists(filename):
|
||
|
|
return
|
||
|
|
|
||
|
|
with open(filename, "r", newline='', encoding="utf8") as file:
|
||
|
|
reader = csv.DictReader(file)
|
||
|
|
|
||
|
|
for row in reader:
|
||
|
|
artist = Artist(row["artist"], float(row["score"]), row["category"])
|
||
|
|
self.artists.append(artist)
|
||
|
|
self.cats.add(artist.category)
|
||
|
|
|
||
|
|
def categories(self):
|
||
|
|
return sorted(self.cats)
|