66 lines
2.4 KiB
Python
66 lines
2.4 KiB
Python
import argparse
|
|
import os
|
|
from random import shuffle
|
|
|
|
import numpy as np
|
|
from matplotlib import pyplot as plt
|
|
from PIL import Image
|
|
|
|
parser = argparse.ArgumentParser(description="Inspect image versions")
|
|
parser.add_argument("-f", "--base_folder", type=str, required=True, help="Base folder to inspect")
|
|
parser.add_argument("-opt_fg_ratio", type=float, default=0.3, help="Optimal foreground ratio")
|
|
|
|
args = parser.parse_args()
|
|
|
|
bg_folder = os.path.join(args.base_folder, "backgrounds")
|
|
fg_folder = os.path.join(args.base_folder, "foregrounds")
|
|
|
|
classes = os.listdir(fg_folder)
|
|
classes = sorted(classes, key=lambda x: int(x[1:]))
|
|
assert len(classes) in [200, 1_000], f"Expected 200 or 1_000 classes, got {len(classes)}"
|
|
|
|
total_images = set()
|
|
for in_cls in classes:
|
|
cls_images = {
|
|
os.path.join(in_cls, "_".join(img.split(".")[0].split("_")[:3]))
|
|
for img in os.listdir(os.path.join(fg_folder, in_cls))
|
|
}
|
|
total_images.update(cls_images)
|
|
total_images = list(total_images)
|
|
shuffle(total_images)
|
|
print(total_images[:5], "...")
|
|
|
|
# in_cls to print name/lemma
|
|
with open("wordnet_data/tinyimagenet_synset_names.txt", "r") as f:
|
|
in_cls_to_name = {line.split(":")[0].strip(): line.split(":")[1].strip() for line in f.readlines() if len(line) > 2}
|
|
|
|
for image_name in total_images:
|
|
in_cls, img_name = image_name.split("/")
|
|
versions = set()
|
|
for img in os.listdir(os.path.join(fg_folder, in_cls)):
|
|
if img.startswith(img_name):
|
|
versions.add(img)
|
|
if len(versions) <= 1:
|
|
print(f"Image {image_name} has only one version")
|
|
continue
|
|
versions = sorted(list(versions))
|
|
|
|
bg_img = Image.open(os.path.join(bg_folder, in_cls, f"{versions[0].split('.')[0]}.JPEG"))
|
|
|
|
# plot all versions at once
|
|
fig, axs = plt.subplots(1, len(versions) + 1, figsize=(15, 5))
|
|
axs[0].imshow(bg_img)
|
|
axs[0].axis("off")
|
|
axs[0].set_title("Background")
|
|
for version, ax in zip(versions, axs[1:]):
|
|
img = Image.open(os.path.join(fg_folder, in_cls, version))
|
|
img_mask = np.array(img.convert("RGBA").split()[-1])
|
|
fg_dens_fact = np.sum(img_mask) / (255 * img_mask.size)
|
|
fg_ratio = np.sum(img_mask) / (255 * bg_img.size[0] * bg_img.size[1])
|
|
ax.imshow(img)
|
|
ax.axis("off")
|
|
ax.set_title(f"{version}\nfg_rat: {fg_ratio:.2f} dist to optimal: {abs(fg_ratio - args.opt_fg_ratio):.2f}")
|
|
fig.suptitle(in_cls_to_name[in_cls])
|
|
plt.show()
|
|
plt.close()
|