Upload app.py

app.py

@@ -0,0 +1,73 @@
+from transformers import  ViTModel, ViTImageProcessor
+from PIL import Image, ImageOps
+import gradio as gr
+import torch
+from datasets import Dataset
+from torch.nn import CosineSimilarity
+
+image_processor = ViTImageProcessor.from_pretrained("vit-base-patch16-224")
+image_encoder = ViTModel.from_pretrained("output/image_encoder/epoch_29").eval().to("cuda")
+scribble_encoder = ViTModel.from_pretrained("output/scibble_encoder/epoch_29").eval().to("cuda")
+
+candidates: Dataset = None
+
+cosinesimilarity = CosineSimilarity()
+
+
+
+def load_candidates(candidate_dir):
+    def preprocess(examples):
+        images = [image.convert("RGB") for image in examples["image"]]
+        examples["image_embedding"] = image_encoder(image_processor(images, return_tensors="pt")["pixel_values"].to("cuda"))["pooler_output"]
+        return examples
+
+    dataset = [dict(image=Image.open(tempfile.name).convert("RGB").resize((224, 224))) for tempfile in candidate_dir]
+    dataset = Dataset.from_list(dataset)
+    with torch.no_grad():
+        dataset = dataset.map(preprocess, batched=True, batch_size=1024)
+
+    return dataset
+
+
+def load_candidates_in_cache(candidate_files):
+    global candidates
+    candidates = load_candidates(candidate_files)
+
+
+def scribble_matching(input_img: Image):
+    input_img = ImageOps.invert(input_img)
+
+    scribble = input_img
+    scribble_embedding = scribble_encoder(image_processor(scribble, return_tensors="pt")["pixel_values"].to("cuda"))["pooler_output"].to("cpu")
+    image_embeddings = torch.tensor(candidates["image_embedding"], dtype=torch.float32)
+
+
+    
+    sim = cosinesimilarity(scribble_embedding, image_embeddings)
+
+    predicts = torch.topk(sim, k=15)
+
+    output_imgs = candidates[predicts.indices.tolist()]["image"]
+    labels = predicts.values.tolist()
+    labels = [f"{label:.3f}" for label in labels]
+
+    return list(zip([input_img] + output_imgs, ["preview"] + labels))
+
+
+def main():
+    with gr.Blocks() as demo:
+        with gr.Row():
+            input_img = gr.Image(type="pil", label="scribble", height=512, width=512, source="canvas", tool="color-sketch", brush_radius=10)
+            prediction_gallery = gr.Gallery(min_width=512, columns=4, show_label=True, )
+
+        with gr.Row():
+            candidate_dir = gr.File(file_count="directory", min_width=300, height=300)
+            load_candidates_btn = gr.Button("Load", variant="secondary", size="sm")
+        btn = gr.Button("Scribble Matching", variant="primary")
+        load_candidates_btn.click(fn=load_candidates_in_cache, inputs=[candidate_dir])
+        btn.click(fn=scribble_matching, inputs=[input_img], outputs=[prediction_gallery])
+
+    demo.launch(debug=True)
+
+if __name__ == "__main__":
+    main()