test

app.py

@@ -111,16 +111,16 @@ def draw_optvq(data, code):
     image = Image.open(buf)
     return image
 
-# def draw_process(x, y, std):
-#     data = torch.randn(N_data, dim)
-#     code = torch.randn(N_code, dim) * std
-#     code[:, 0] += x
-#     code[:, 1] += y
+def draw_process(x, y, std):
+    data = torch.randn(N_data, dim)
+    code = torch.randn(N_code, dim) * std
+    code[:, 0] += x
+    code[:, 1] += y
 
-#     image_NN = draw_NN(data, code)
-#     image_optvq = draw_optvq(data, code)
+    image_NN = draw_NN(data, code)
+    image_optvq = draw_optvq(data, code)
     
-#     return image_NN, image_optvq
+    return image_NN, image_optvq
 
 class Handler:
     def __init__(self, device):
@@ -172,14 +172,14 @@ class Handler:
             quant, *_ = self.vqgan.encode(img)
             vqgan_rec = self.vqgan.decode(quant)
             # revq
-            # lat = self.vae.encode(img).latent
-            # lat = lat.contiguous()
-            # lat = self.preprocesser(lat)
-            # lat = self.revq.quantize(lat)
-            # revq_rec = self.revq.decode(lat)
-            # revq_rec = revq_rec.contiguous()
-            # revq_rec = self.preprocesser.inverse(revq_rec)
-            # revq_rec = self.vae.decode(revq_rec).sample
+            lat = self.vae.encode(img).latent
+            lat = lat.contiguous()
+            lat = self.preprocesser(lat)
+            lat = self.revq.quantize(lat)
+            revq_rec = self.revq.decode(lat)
+            revq_rec = revq_rec.contiguous()
+            revq_rec = self.preprocesser.inverse(revq_rec)
+            revq_rec = self.vae.decode(revq_rec).sample
             # optvq
             quant, *_ = self.optvq.encode(img)
             optvq_rec = self.optvq.decode(quant)
@@ -188,92 +188,106 @@ class Handler:
         img = self.tensor_to_image(img)
         basevq_rec = self.tensor_to_image(basevq_rec)
         vqgan_rec = self.tensor_to_image(vqgan_rec)
-        # revq_rec = self.tensor_to_image(revq_rec)
+        revq_rec = self.tensor_to_image(revq_rec)
         optvq_rec = self.tensor_to_image(optvq_rec)
         # print("Shapes:", img.shape, basevq_rec.shape, vqgan_rec.shape, revq_rec.shape)
-        # return img, basevq_rec, vqgan_rec, revq_rec
-        return basevq_rec, vqgan_rec, optvq_rec
+        return basevq_rec, vqgan_rec, revq_rec
+        # return basevq_rec, vqgan_rec, optvq_rec
 
-def draw_process(x, y, std):
-    img = (np.random.rand(256, 256, 3) * 255).astype(np.uint8)
-    return img, img
-
-demo2 = gr.Interface(
-    fn=draw_process,
-    inputs=[
-        gr.Slider(label="x", value=0, minimum=-10, maximum=10, step=0.1),
-        gr.Slider(label="y", value=0, minimum=-10, maximum=10, step=0.1),
-        gr.Slider(label="std", value=1, minimum=0, maximum=5, step=0.1)
-    ],
-    outputs=[
-        gr.Image(label="NN", type="numpy"),
-        gr.Image(label="OptVQ", type="numpy")
-    ],
-    title="Demo 2: 2D Matching Visualization",
-    description="Visualize nearest neighbor vs. optimal transport matching for synthetic 2D data."
-)
-
-demo2.launch()
-
-# if __name__ == "__main__":
-#     # create the model handler
-#     # handler = Handler(device=device)
-
-#     print("Creating Gradio interface...")
-
-#     demo2 = gr.Interface(
+# def draw_process(x, y, std):
+#     img = (np.random.rand(256, 256, 3) * 255).astype(np.uint8)
+#     return img, img
+
+# demo2 = gr.Interface(
 #     fn=draw_process,
-#         inputs=[
-#             gr.Slider(label="x", value=0, minimum=-10, maximum=10, step=0.1),
-#             gr.Slider(label="y", value=0, minimum=-10, maximum=10, step=0.1),
-#             gr.Slider(label="std", value=1, minimum=0, maximum=5, step=0.1)
-#         ],
-#         outputs=[
-#             gr.Image(label="NN", type="numpy"),
-#             gr.Image(label="OptVQ", type="numpy")
-#         ],
-#         title="Demo 2: 2D Matching Visualization",
-#         description="Visualize nearest neighbor vs. optimal transport matching for synthetic 2D data."
-#     )
-
-#     # 合并两个 interface 成 Tabbed UI
-#     demo = gr.TabbedInterface(
-#         interface_list=[demo2],
-#         tab_names=["2D Matching"]
-#     )
-
-#     demo.launch()
-
-#     # create the interface
-#     # with gr.Blocks() as demo:
-#         # gr.Textbox(value="This demo shows the image reconstruction comparison between ReVQ and other methods. The input image is resized to 256 x 256 and then fed into the models. The output images are the reconstructed images from the latent codes.", label="Demo 1: Image reconstruction results")
-#         # with gr.Row():
-#         #     with gr.Column():
-#         #         image_input = gr.Image(label="Input data", image_mode="RGB", type="numpy")
-#         #         btn_demo1 = gr.Button(value="Run reconstruction")
-#         #     image_basevq = gr.Image(label="BaseVQ rec.")
-#         #     image_vqgan = gr.Image(label="VQGAN rec.")
-#         #     image_revq = gr.Image(label="ReVQ rec.")
-#         # btn_demo1.click(fn=handler.process_image, inputs=[image_input], outputs=[image_basevq, image_vqgan, image_revq])
-
-#         # gr.Textbox(value="This demo shows the 2D visualizations of nearest neighbor and optimal transport (OptVQ) methods. The data points are randomly generated from a normal distribution, and the matching results are shown as arrows with different colors.", label="Demo 2: 2D visualizations of matching results")
-#         # gr.Markdown("### Demo 2: 2D visualizations of matching results\n"
-#         #         "This demo shows the 2D visualizations of nearest neighbor and optimal transport (OptVQ) methods. "
-#         #         "The data points are randomly generated from a normal distribution, and the matching results are shown as arrows with different colors.")
-#         # with gr.Row():
-#         #     with gr.Column():
-#         #         input_x = gr.Slider(label="x", value=0, minimum=-10, maximum=10, step=0.1)
-#         #         input_y = gr.Slider(label="y", value=0, minimum=-10, maximum=10, step=0.1)
-#         #         input_std = gr.Slider(label="std", value=1, minimum=0, maximum=5, step=0.1)
-#         #         btn_demo2 = gr.Button(value="Run 2D example")
-#         #     output_nn = gr.Image(label="NN", interactive=False, type="numpy")
-#         #     output_optvq = gr.Image(label="OptVQ", interactive=False, type="numpy")
+#     inputs=[
+#         gr.Slider(label="x", value=0, minimum=-10, maximum=10, step=0.1),
+#         gr.Slider(label="y", value=0, minimum=-10, maximum=10, step=0.1),
+#         gr.Slider(label="std", value=1, minimum=0, maximum=5, step=0.1)
+#     ],
+#     outputs=[
+#         gr.Image(label="NN", type="numpy"),
+#         gr.Image(label="OptVQ", type="numpy")
+#     ],
+#     title="Demo 2: 2D Matching Visualization",
+#     description="Visualize nearest neighbor vs. optimal transport matching for synthetic 2D data."
+# )
+
+# demo2.launch()
+
+if __name__ == "__main__":
+    # create the model handler
+    handler = Handler(device=device)
+
+    print("Creating Gradio interface...")
+
+    # Demo 1 接口：图像重建
+    demo1 = gr.Interface(
+        fn=handler.process_image,
+        inputs=gr.Image(label="Input Image", type="numpy"),
+        outputs=[
+            gr.Image(label="BaseVQ Reconstruction", type="numpy"),
+            gr.Image(label="VQGAN Reconstruction", type="numpy"),
+            gr.Image(label="ReVQ Reconstruction", type="numpy"),
+            # 若启用 ReVQ：gr.Image(label="ReVQ Reconstruction", type="numpy"),
+        ],
+        title="Demo 1: Image Reconstruction",
+        description="Upload an image to see how different VQ models (BaseVQ, VQGAN, ReVQ) reconstruct it from latent codes."
+    )
+
+    demo2 = gr.Interface(
+        fn=draw_process,
+        inputs=[
+            gr.Slider(label="x", value=0, minimum=-10, maximum=10, step=0.1),
+            gr.Slider(label="y", value=0, minimum=-10, maximum=10, step=0.1),
+            gr.Slider(label="std", value=1, minimum=0, maximum=5, step=0.1)
+        ],
+        outputs=[
+            gr.Image(label="NN", type="numpy"),
+            gr.Image(label="OptVQ", type="numpy")
+        ],
+        title="Demo 2: 2D Matching Visualization",
+        description="Visualize nearest neighbor vs. optimal transport matching for synthetic 2D data."
+    )
+
+    # 合并两个 interface 成 Tabbed UI
+    demo = gr.TabbedInterface(
+        interface_list=[demo1, demo2],
+        tab_names=["Image Reconstruction", "2D Matching"]
+    )
+
+    demo.launch(share=True)
+
+    # create the interface
+    # with gr.Blocks() as demo:
+        # gr.Textbox(value="This demo shows the image reconstruction comparison between ReVQ and other methods. The input image is resized to 256 x 256 and then fed into the models. The output images are the reconstructed images from the latent codes.", label="Demo 1: Image reconstruction results")
+        # with gr.Row():
+        #     with gr.Column():
+        #         image_input = gr.Image(label="Input data", image_mode="RGB", type="numpy")
+        #         btn_demo1 = gr.Button(value="Run reconstruction")
+        #     image_basevq = gr.Image(label="BaseVQ rec.")
+        #     image_vqgan = gr.Image(label="VQGAN rec.")
+        #     image_revq = gr.Image(label="ReVQ rec.")
+        # btn_demo1.click(fn=handler.process_image, inputs=[image_input], outputs=[image_basevq, image_vqgan, image_revq])
+
+        # gr.Textbox(value="This demo shows the 2D visualizations of nearest neighbor and optimal transport (OptVQ) methods. The data points are randomly generated from a normal distribution, and the matching results are shown as arrows with different colors.", label="Demo 2: 2D visualizations of matching results")
+        # gr.Markdown("### Demo 2: 2D visualizations of matching results\n"
+        #         "This demo shows the 2D visualizations of nearest neighbor and optimal transport (OptVQ) methods. "
+        #         "The data points are randomly generated from a normal distribution, and the matching results are shown as arrows with different colors.")
+        # with gr.Row():
+        #     with gr.Column():
+        #         input_x = gr.Slider(label="x", value=0, minimum=-10, maximum=10, step=0.1)
+        #         input_y = gr.Slider(label="y", value=0, minimum=-10, maximum=10, step=0.1)
+        #         input_std = gr.Slider(label="std", value=1, minimum=0, maximum=5, step=0.1)
+        #         btn_demo2 = gr.Button(value="Run 2D example")
+        #     output_nn = gr.Image(label="NN", interactive=False, type="numpy")
+        #     output_optvq = gr.Image(label="OptVQ", interactive=False, type="numpy")
         
-#         # # set the function
-#         # input_x.change(fn=draw_process, inputs=[input_x, input_y, input_std], outputs=[output_nn, output_optvq])
-#         # input_y.change(fn=draw_process, inputs=[input_x, input_y, input_std], outputs=[output_nn, output_optvq])
-#         # input_std.change(fn=draw_process, inputs=[input_x, input_y, input_std], outputs=[output_nn, output_optvq])
-#         # btn_demo2.click(fn=draw_process, inputs=[input_x, input_y, input_std], outputs=[output_nn, output_optvq])
-#         # btn_demo2.click(fn=draw_process, inputs=[input_x, input_y, input_std], outputs=[output_nn, output_optvq])
-
-#     # demo.launch()
+        # # set the function
+        # input_x.change(fn=draw_process, inputs=[input_x, input_y, input_std], outputs=[output_nn, output_optvq])
+        # input_y.change(fn=draw_process, inputs=[input_x, input_y, input_std], outputs=[output_nn, output_optvq])
+        # input_std.change(fn=draw_process, inputs=[input_x, input_y, input_std], outputs=[output_nn, output_optvq])
+        # btn_demo2.click(fn=draw_process, inputs=[input_x, input_y, input_std], outputs=[output_nn, output_optvq])
+        # btn_demo2.click(fn=draw_process, inputs=[input_x, input_y, input_std], outputs=[output_nn, output_optvq])
+
+    # demo.launch()