app.py

import gradio as gr
from transformers import pipeline
from datasets import load_dataset
import requests
import traceback
import json
import os

dataset = load_dataset("SaladSlayer00/twin_matcher_data")

image_classifier = pipeline("image-classification", model="SaladSlayer00/twin_matcher_beta")

def format_info(info_json):

        
    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



def fetch_info(celebrity_label):
        
    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.getenv('TOKEN')
    response = requests.get(api_url, headers={'X-Api-Key': token})
    if response.status_code == 200:
        return format_info(response.text)
    else:
        return "A shining star for sure."

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


def predict_and_fetch_images(input_image):

        # Use the image classifier pipeline
    predictions = image_classifier(input_image)
    top_prediction = max(predictions, key=lambda x: x['score'])
    label, score = top_prediction['label'], top_prediction['score']

        # Fetch images for the predicted label
    images = fetch_images_for_label(label)

        # Fetch information for the predicted label
    info = fetch_info(label)

    return label, score, images, info, "No Error"


example_images = [
    "images/megan_fox.png",
    "images/chris_evans.png",
    "images/millie_bobby_brown.png",
    "images/alvaro_morte.png",
    "images/amber_heard.png"
]
        
# Gradio interface
iface = gr.Interface(
    fn=predict_and_fetch_images,
    inputs=gr.Image(type="pil", label="Upload or Take a Snapshot"),
    outputs=[
        "text",  # Predicted label
        "number",  # Prediction score
        gr.Gallery(label="Lookalike Images"),  # Slideshow component for images
        "html",  # Info/Description as HTML
        gr.Textbox(type="text", label="Feedback", placeholder="Provide feedback here")  # Feedback textbox
    ],
    examples=example_images,
    live=True,
    title="Celebrity Lookalike Predictor",
    description="Take a snapshot or upload an image to see which celebrity you look like!"
)


iface.launch()