Update app.py

app.py

@@ -1,36 +1,50 @@
-import gradio as gr
-import torch
-import numpy as np
-from PIL import Image
+import gradio as gr  # Import Gradio library for creating UI
+import torch  # Import PyTorch for YOLOv5 model
+from PIL import Image  # Import Pillow for image manipulation
 
-# Load the custom-trained YOLOv5 model
+# Load YOLOv5 model from Ultralytics' GitHub repository
 model = torch.hub.load('ultralytics/yolov5', 'custom', path='best.pt', _verbose=False)
 
-def gradio_wrapper(img):
-    global model
-    # Run the model on the input image and get the results
-    results = model(img)
-    # Render the results and return the annotated image
-    return results.render()[0]
+# Define a function to predict objects in an image
+def predict(image_in_img, image_in_video):
+    global model  # Use the global YOLOv5 model
+    if image_in_video == None and image_in_img == None:  # If both inputs are None, raise an error
+        raise gr.Error("Please upload an image.")
+    if image_in_video or image_in_img:  # If either input is not None,
+        image = image_in_video or image_in_img  # set the image variable to the non-None input
+        return model(image).render()[0]  # Use the YOLOv5 model to predict objects in the image and return the rendered output
 
-# Set up the Gradio image input component with the option to upload an image or use the webcam
-s = gr.Radio(["upload", "webcam"], type="index")
+# Define a function to toggle between webcam and file inputs
+def toggle(choice):
+    if choice == "webcam":  # If "webcam" is selected,
+        return gr.update(visible=True, value=None), gr.update(visible=False, value=None)  # Show webcam input and hide file input
+    else:  # Otherwise, if "file" is selected,
+        return gr.update(visible=False, value=None), gr.update(visible=True, value=None)  # Show file input and hide webcam input
 
-if s == "upload":
-    image = gr.inputs.Image(source="upload")
-else:
-    image = gr.inputs.Image(source="webcam")
+# Create Gradio UI blocks
+with gr.Blocks() as blocks:
+    gr.Markdown("WebCam or Upload?""")  # Display text in Markdown format
+    with gr.Row():  # Create a row of UI elements
+        with gr.Column():  # Create a column of UI elements
+            # Create a radio button to choose between webcam and file inputs
+            image_or_file_opt = gr.Radio(["webcam", "file"], value="webcam",
+                                         label="How would you like to upload your image?")
+            # Create an image input for the webcam
+            image_in_video = gr.Image(source="webcam", type="filepath")
+            # Create an image input for a file, initially hidden
+            image_in_img = gr.Image(
+                source="upload", visible=False, type="filepath")
 
+            # Bind the toggle function to the radio button to switch between webcam and file inputs
+            image_or_file_opt.change(fn=toggle, inputs=[image_or_file_opt],
+                                     outputs=[image_in_video, image_in_img], queue=False)
+        with gr.Column():  # Create another column of UI elements
+            # Create an output image to display the predicted objects
+            image_out = gr.Image()
+    # Create a button to run the prediction function and display the output image
+    run_btn = gr.Button("Run")
+    run_btn.click(fn=predict, inputs=[
+                  image_in_img, image_in_video], outputs=[image_out])
 
-# Create the Gradio interface with the gradio_wrapper function, the image input component, and an image output component
-demo = gr.Interface(
-    gradio_wrapper,
-    image,
-    'image',
-    live=True,
-    title="CiclopeIA",
-    description="App based on the CiclopeIA project from Saturdays AI. Identifies the value of Euro banknotes."
-)
-
-# Launch the Gradio interface
-demo.launch()
+# Launch the Gradio UI blocks
+blocks.launch()