Update app.py

app.py

@@ -1,25 +1,33 @@
 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_video(video):
-    # Classify the uploaded video and return the results
-    predictions = video_model(video)
-    
-    # Create a dictionary of labels and their corresponding scores
-    result = {pred["label"]: pred["score"] for pred in predictions}
-    
-    # Return the result dictionary
-    return result
+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_video,
-    inputs=gr.Video(label="Upload Video"),
-    outputs=gr.Label(num_top_classes=3, label="Predictions"),
-    title="Video Classification App",
-    description="Upload a video to classify its content."
+    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