import gradio as gr
from transformers import pipeline

# Load the models for image and video classification
#image_model = pipeline("image-classification", model="Rohit1412/Deepfake", device=0)
video_model = pipeline("video-classification", model="Rohit1412/deepfakerohit2.0")

def classify_input(input_data):
    # Check if the input is an image or a video
    if isinstance(input_data, str) and input_data.endswith(('.mp4', '.mov', '.avi')):
        # Classify the uploaded video
        predictions = video_model(input_data)
        # Create a dictionary of labels and their corresponding scores
        result = {pred["label"]: pred["score"] for pred in predictions}
        return result
    else:
        # Classify the uploaded image
        predictions = image_model(input_data)
        return predictions[0]['label'], predictions[0]['score']

# Create Gradio interface
interface = gr.Interface(
    fn=classify_input,
    inputs=gr.inputs.Audio(source="upload", type="filepath", label="Upload Image or Video"),
    outputs=[
        gr.Label(num_top_classes=3, label="Predictions"),
        gr.Textbox(label="Confidence", visible=False)  # Hide confidence for video output
    ],
    title="Image and Video Classification App",
    description="Upload an image or a video to classify its content. Note: The model is still under training; it may give incorrect outputs."
)

# Launch the interface
if __name__ == "__main__":
    interface.launch(debug=True)