Update app.py

app.py

@@ -1,27 +1,44 @@
-import gradio as gr
 from transformers import pipeline
+import gradio as gr
 
+# Load the models
+image_model = pipeline("image-classification", model="Rohit1412/Deepfake", device=0)
 video_model = pipeline("video-classification", model="Rohit1412/deepfakerohit2.0")
 
+def classify_image(image):
+    try:
+        predictions = image_model(image)
+        result = {pred["label"]: pred["score"] for pred in predictions}
+        return result
+    except Exception as e:
+        return {"error": f"An error occurred during image classification Make sure you choose the correct model: {str(e)}"}
+
 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}
+    try:
+        predictions = video_model(video)
+        result = {pred["label"]: pred["score"] for pred in predictions}
+        return result
+    except Exception as e:
+        return {"error": f"An error occurred during video classification Make sure you choose the correct model: {str(e)}"}
+
+# Create Gradio Blocks interface
+with gr.Blocks() as app:
+    gr.Markdown("# Video and Image Classification App")
     
-    # Return the result dictionary
-    return result
+    with gr.Tab("Image Classification"):
+        image_input = gr.Image(label="Upload Image")
+        image_output = gr.Label(num_top_classes=3, label="Predictions")
+        image_button = gr.Button("Classify Image")
+        
+        image_button.click(fn=classify_image, inputs=image_input, outputs=image_output)
 
-# 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."
-)
+    with gr.Tab("Video Classification"):
+        video_input = gr.Video(label="Upload Video")
+        video_output = gr.Label(num_top_classes=3, label="Predictions")
+        video_button = gr.Button("Classify Video")
+        
+        video_button.click(fn=classify_video, inputs=video_input, outputs=video_output)
 
 # Launch the interface
 if __name__ == "__main__":
-    interface.launch(debug=True)
+    app.launch(debug=True)