Update app.py

app.py

@@ -7,6 +7,7 @@ from torchvision.transforms import (
     ToTensor,
     Normalize,
     InterpolationMode,
+    CenterCrop,
 )
 from PIL import Image
 import gradio as gr
@@ -15,7 +16,7 @@ print("starting...")
 (ys,) = np.load("embs.npz").values()
 print("loaded embs")
 model = torch.load(
-    "style-extractor-v0.2.0.ckpt",
+    "style-extractor-v0.3.0.ckpt",
     map_location="cpu",
 )
 print("loaded extractor")
@@ -27,15 +28,40 @@ d = ys.shape[1]
 index = faiss.IndexHNSWFlat(d, 32)
 print("building index")
 index.add(ys)
-print('index built')
+print("index built")
+
+
+def MyResize(area, d):
+    def f(im: Image):
+        w, h = im.size
+        s = (area / w / h) ** 0.5
+        wd, hd = int(s * w / d), int(s * h / d)
+        e = lambda a, b: 1 - min(a, b) / max(a, b)
+        wd, hd = min(
+            (
+                (ww * d, hh * d)
+                for ww, hh in [(wd + i, hd + j) for i in (0, 1) for j in (0, 1)]
+                if ww * d * hh * d <= area
+            ),
+            key=lambda wh: e(wh[0] / wh[1], w / h),
+        )
+
+        return Compose(
+            [
+                Resize(
+                    (int(h * wd / w), wd) if wd / w > hd / h else (hd, int(w * hd / h)),
+                    InterpolationMode.BICUBIC,
+                ),
+                CenterCrop((hd, wd)),
+            ]
+        )(im)
+
+    return f
+
+
 tf = Compose(
     [
-        Resize(
-            size=336,
-            interpolation=InterpolationMode.BICUBIC,
-            max_size=None,
-            antialias=True,
-        ),
+        MyResize((518 * 1.3) ** 2, 14),
         ToTensor(),
         Normalize(mean=[0.4850, 0.4560, 0.4060], std=[0.2290, 0.2240, 0.2250]),
     ]
@@ -56,6 +82,7 @@ def f(im):
     D, I = index.search(get_emb(im), n_outputs)
     return [f"Distance: {d:.1f}\n![]({urls[i]})" for d, i in zip(D[0], I[0])]
 
+
 print("preparing gradio")
 with gr.Blocks() as demo:
     gr.Markdown(
@@ -70,4 +97,4 @@ with gr.Blocks() as demo:
                 outputs.append(gr.Markdown(label=f"#{len(outputs) + 1}"))
     btn.click(f, img, outputs)
 print("starting gradio")
-demo.launch()
+demo.launch()