select_bucket
A select_bucket
függvény felelős egy bemeneti kép megfelelő felbontásának (méretének) meghatározásáért, előre definiált felbontások és képarányok, valamint a kép számára megengedett maximális terület alapján. A függvény két argumentumot vesz át: image_width
és image_height
, amelyek a bemeneti kép méreteit reprezentálják.
Íme a függvény logikájának részletes magyarázata:
Kiszámítja a bemeneti kép képarányát az
image_width
ésimage_height
hányadosaként.Ha a
no_upscale
értéke False, a függvény megpróbál találni egy olyan előre definiált felbontást, amely megegyezik vagy közel áll a bemeneti kép képarányához. Így működik:a. Ha a bemeneti kép felbontása már szerepel az előre definiált felbontások halmazában (
self.predefined_resos_set
), változatlanul hagyja a felbontást.b. Ha nem, kiszámítja a különbséget a bemeneti képarány és az összes előre definiált képarány (
self.predefined_aspect_ratios
) között. Ezután kiválasztja azt az előre definiált felbontást, amelynek a legkisebb a képarány különbsége (predefined_bucket_id
).c. A kiválasztott előre definiált felbontás (
reso
) alapján meghatározza a méretezési tényezőt a bemeneti kép átméretezéséhez. Ha a bemeneti képarány nagyobb, mint az előre definiált képarány, az előre definiált felbontás magasságához igazítja. Ellenkező esetben a szélességhez igazítja.d. Az átméretezett méretek (
resized_size
) kiszámítása a bemeneti méretek és a méretezési tényező szorzatának legközelebbi egészre kerekítésével történik.Ha a
no_upscale
értéke True, a függvény csak a kép csökkentését (lekicsinyítését) végzi:a. Ha a bemeneti kép területe (
image_width * image_height
) meghaladja a megengedett maximális területet (self.max_area
), kiszámítja az átméretezett méreteket (resized_width
ésresized_height
) úgy, hogy megőrzi a képarányt, miközben biztosítja, hogy a terület nem haladja meg aself.max_area
értéket.b. Ezután kerekíti az átméretezett méreteket a
self.reso_steps
(előre definiált felbontási lépésköz) legközelebbi többszörösére. Ez úgy történik, hogy azt a kerekítést választja, amely a legközelebb áll az eredeti képarányhoz.c. A végső átméretezett méretek (
resized_size
) a kerekített értékekre állítódnak be.d. Ha a bemeneti kép területe nem haladja meg a
self.max_area
értéket, a függvény aresized_size
értékét az eredeti bemeneti méretekre állítja, mivel nincs szükség átméretezésre.e. A függvény ezután kiszámítja a bucket felbontást (
reso
) úgy, hogy megkeresi a legnagyobb olyan méreteket, amelyek aself.reso_steps
többszörösei és kisebbek vagy egyenlőek az átméretezett méretekkel.A függvény hozzáadja a kiválasztott felbontást (
reso
) az egyedi felbontások halmazához (self.add_if_new_reso(reso)
).Végül kiszámítja a kiválasztott felbontás és az eredeti képarány közötti hibát (
ar_error
), és visszaadja a(reso, resized_size, ar_error)
tuple-t.
A get_crop_ltrb
függvény, amely egy statikus metódus, kiszámítja a bal, felső, jobb és alsó koordinátákat egy kép kivágásához, hogy illeszkedjen egy adott bucket felbontáshoz. Két argumentumot vesz át: bucket_reso
(a célzott felbontás) és image_size
(az eredeti képméret). A függvény meghatározza, hogy a bucket felbontás magasságához vagy szélességéhez igazodjon-e a képarányok alapján, kiszámítja az átméretezett méreteket, majd kiszámítja a kivágási koordinátákat a bucket felbontás és az átméretezett méretek közötti különbség alapján.
Forráskód
def select_bucket(self, image_width, image_height):
aspect_ratio = image_width / image_height
if not self.no_upscale:
# Nagyítás és kicsinyítés végrehajtása
# Prioritást ad az azonos felbontásnak, mert lehet ugyanaz a képarány
# (finomhangolás esetén no_upscale=True előfeldolgozással)
reso = (image_width, image_height)
if reso in self.predefined_resos_set:
pass
else:
ar_errors = self.predefined_aspect_ratios - aspect_ratio
predefined_bucket_id = np.abs(ar_errors).argmin() # A felbontáson kívül a legkisebb képarány hibával rendelkező
reso = self.predefined_resos[predefined_bucket_id]
ar_reso = reso[0] / reso[1]
if aspect_ratio > ar_reso: # Ha a szélesség hosszabb, a magassághoz igazítjuk
scale = reso[1] / image_height
else:
scale = reso[0] / image_width
resized_size = (int(image_width * scale + 0.5), int(image_height * scale + 0.5))
# logger.info(f"use predef, {image_width}, {image_height}, {reso}, {resized_size}")
else:
# Csak kicsinyítés végrehajtása
if image_width * image_height > self.max_area:
# Ha a kép túl nagy, a bucket-et úgy határozzuk meg, hogy csökkentjük a képarányt megtartva
resized_width = math.sqrt(self.max_area * aspect_ratio)
resized_height = self.max_area / resized_width
assert abs(resized_width / resized_height - aspect_ratio) < 1e-2, "a képarány illegális"
# Az átméretezés után a rövid vagy hosszú oldalt a reso_steps többszörösévé tesszük: azt választjuk, amelyiknek kisebb a képarány különbsége
# Ugyanaz a logika, mint az eredeti bucketingnél
b_width_rounded = self.round_to_steps(resized_width)
b_height_in_wr = self.round_to_steps(b_width_rounded / aspect_ratio)
ar_width_rounded = b_width_rounded / b_height_in_wr
b_height_rounded = self.round_to_steps(resized_height)
b_width_in_hr = self.round_to_steps(b_height_rounded * aspect_ratio)
ar_height_rounded = b_width_in_hr / b_height_rounded
# logger.info(b_width_rounded, b_height_in_wr, ar_width_rounded)
# logger.info(b_width_in_hr, b_height_rounded, ar_height_rounded)
if abs(ar_width_rounded - aspect_ratio) < abs(ar_height_rounded - aspect_ratio):
resized_size = (b_width_rounded, int(b_width_rounded / aspect_ratio + 0.5))
else:
resized_size = (int(b_height_rounded * aspect_ratio + 0.5), b_height_rounded)
# logger.info(resized_size)
else:
resized_size = (image_width, image_height) # Nincs szükség átméretezésre
# A kép méretét a bucket méreténél kisebbé tesszük (kivágás kitöltés nélkül)
bucket_width = resized_size[0] - resized_size[0] % self.reso_steps
bucket_height = resized_size[1] - resized_size[1] % self.reso_steps
# logger.info(f"use arbitrary {image_width}, {image_height}, {resized_size}, {bucket_width}, {bucket_height}")
reso = (bucket_width, bucket_height)
self.add_if_new_reso(reso)
ar_error = (reso[0] / reso[1]) - aspect_ratio
return reso, resized_size, ar_error
@staticmethod
def get_crop_ltrb(bucket_reso: Tuple[int, int], image_size: Tuple[int, int]):
# A bal/felső kivágás kiszámítása a Stability AI előfeldolgozása szerint. A jobb oldali kivágás a tükrözési augmentációhoz számítódik.
bucket_ar = bucket_reso[0] / bucket_reso[1]
image_ar = image_size[0] / image_size[1]
if bucket_ar > image_ar:
# Ha a bucket szélesebb, a magassághoz igazítjuk
resized_width = bucket_reso[1] * image_ar
resized_height = bucket_reso[1]
else:
resized_width = bucket_reso[0]
resized_height = bucket_reso[0] / image_ar
crop_left = (bucket_reso[0] - resized_width) // 2
crop_top = (bucket_reso[1] - resized_height) // 2
crop_right = crop_left + resized_width
crop_bottom = crop_top + resized_height
return crop_left, crop_top, crop_right, crop_bottom
class BucketBatchIndex(NamedTuple):
bucket_index: int
bucket_batch_size: int
batch_index: int