Update app.py

app.py

@@ -1,48 +1,96 @@
 import gradio as gr
+from transformers import pipeline
+from datasets import load_dataset
 import requests
-import io
-import os
-import tempfile
-from PIL import Image
 import traceback
+import json
 
-token = os.getenv('TOKEN')
-# API details
-API_URL = "https://api-inference.huggingface.co/models/SaladSlayer00/twin_matcher"
-headers = {"Authorization": token}
-
-# Function to query the API with an image file
-def query(filename):
-    with open(filename, "rb") as f:
-        data = f.read()
-    response = requests.post(API_URL, headers=headers, data=data)
-    response.raise_for_status()  # Ensure successful response
-    return response.json()
-
-# Function to process the image and predict
-def predict(image):
+dataset = load_dataset("SaladSlayer00/twin_matcher")
+
+image_classifier = pipeline("image-classification", model="SaladSlayer00/twin_matcher")
+
+def format_info(info_json):
     try:
-        with tempfile.NamedTemporaryFile(delete=False, suffix='.jpeg') as tmp_file:
-            image.save(tmp_file, format="JPEG")
-            tmp_file_path = tmp_file.name
+        
+        info_data = json.loads(info_json)
+
+        formatted_info = "<table style='border-collapse: collapse; width: 80%; margin: 20px;'>"
+        formatted_info += "<tr style='background-color: #f2f2f2;'>"
+        for key in info_data[0].keys():
+            formatted_info += f"<th style='border: 1px solid #dddddd; text-align: left; padding: 8px;'><b>{key.capitalize()}</b></th>"
+        formatted_info += "</tr>"
+
+        for entry in info_data:
+            formatted_info += "<tr>"
+            for value in entry.values():
+                formatted_info += f"<td style='border: 1px solid #dddddd; text-align: left; padding: 8px;'>{value}</td>"
+            formatted_info += "</tr>"
+        formatted_info += "</table>"
+        return formatted_info
+    except Exception as e:
+        print(f"Error formatting info: {e}")
+        return "Info not available."
+
+def fetch_info(celebrity_label):
+    try:
+        
+        parts = celebrity_label.split("_")
+        formatted_label = " ".join([part.capitalize() for part in parts])
+
+        
+        api_url = f'https://api.api-ninjas.com/v1/celebrity?name={formatted_label}'
+
+        token = os.environ.get('TOKEN')  
+        response = requests.get(api_url, headers={'X-Api-Key': token})
+        if response.status_code == 200:
+            return format_info(response.text)
+        else:
+            return "Description not available."
+    except Exception as e:
+        print(f"Error fetching information: {e}")
+        traceback.print_exc()
+        return "Description not available."
+
+def fetch_images_for_label(label):
+    label_data = dataset['train'].filter(lambda example: example['label'] == label)
+    images = [example['image'] for example in label_data]
+    return images
 
-        predictions = query(tmp_file_path)
 
+def predict_and_fetch_images(input_image):
+    try:
+        
+        predictions = image_classifier(input_image)
         top_prediction = max(predictions, key=lambda x: x['score'])
-        result = (top_prediction['label'], top_prediction['score'])
+        label, score = top_prediction['label'], top_prediction['score']
+
+        
+        images = fetch_images_for_label(label)
 
-        os.remove(tmp_file_path)
+        
+        info = fetch_info(label)
 
-        return result
+        return label, score, images, info, "No Error"
     except Exception as e:
-        print(f"Exception during prediction: {e}")
+        print(f"Error during prediction: {e}")
         traceback.print_exc()
-        return "Error during prediction", "N/A"
-
-# Gradio interface
-with gr.Interface(fn=predict, 
-                  inputs=gr.Image(type="pil"), 
-                  outputs=["text", "number"],
-                  title="Celebrity Lookalike Predictor",
-                  description="Take a snapshot to see which celebrity you look like!") as demo:
-    demo.launch()
+        return "Error during prediction", 0, [], "N/A", str(e)
+
+
+iface = gr.Interface(
+    fn=predict_and_fetch_images,
+    inputs=gr.Image(type="pil", label="Upload or Take a Snapshot"),
+    outputs=[
+        "text",  
+        "number",  
+        gr.Gallery(label="Lookalike Images"),  
+        "html",  
+        gr.Textbox(type="text", label="Feedback", placeholder="Provide feedback here")  # Feedback textbox
+    ],
+    live=True,
+    title="Celebrity Lookalike Predictor",
+    description="Take a snapshot or upload an image to see which celebrity you look like!"
+)
+
+
+iface.launch()