Add image prediction

app.py

@@ -18,6 +18,21 @@ MODELS = {
 }
 
 
+def get_image(text: str):
+    url = "https://wikisearch.uncool.ai/get_image/"
+    try:
+        payload = {
+            "text": text,
+            "image_width": 400
+        }
+        data = json.dumps(payload)
+        response = requests.request("POST", url, headers=headers, data=data)
+        print(response.content)
+        image = json.loads(response.content.decode("utf-8"))["url"]
+    except:
+        image = ""
+    return image
+
 def query(payload, model_name):
     data = json.dumps(payload)
     # print("model url:", MODELS[model_name]["url"])
@@ -145,7 +160,15 @@ if st.button("Run"):
             result = result[0]["generated_text"]
             st.write(result.replace("\n", "  \n"))
             st.text("Translation")
+            translation = translate(result, "en", "id")
             if lang == "id":
-                st.write(translate(result, "en", "id").replace("\n", "  \n"))
+                st.write(translation.replace("\n", "  \n"))
             else:
                 st.write(translate(result, lang, "id").replace("\n", "  \n"))
+
+            image_cat = "https://media.giphy.com/media/vFKqnCdLPNOKc/giphy.gif"
+            image = get_image(translation)
+            if image is not "":
+                st.image(image, width=400)
+            else:
+                st.image(image_cat, width=400)