feat: set output to iamge

app.py

@@ -36,16 +36,15 @@ def search_index(input_image, k: int):
         _, I = index.search(np.array(embedding[0].reshape(1, -1)), k)
         indices = I[0]
         answer = ""
-        for i, index in enumerate(indices[:3]):
-            answer += index
-            # retrieved_img = dataset["train"][int(index)]["image"]
-    return answer
+        for i, index in enumerate(indices[:1]):
+            retrieved_img = dataset["train"][int(index)]["image"]
+    return retrieved_img 
 
 
 app = gr.Interface(
     search_index,
-    inputs=[gr.Image(), gr.Slider(value=3, minimum=1, step=1)],
-    outputs="text",
+    inputs=gr.Image(),
+    outputs="image",
 )
 
 if __name__ == "__main__":