Update app.py

app.py

@@ -92,13 +92,10 @@ def process_frame(frame, selected_model):
             
             plt.close(fig)
         
-        # Convert frame_rgb back to BGR for video processing
-        frame_bgr = cv2.cvtColor(frame_rgb, cv2.COLOR_RGB2BGR)
-        
-        return frame_rgb, frame_bgr, plot_img
+        return frame_rgb, frame_rgb, plot_img
     except Exception as e:
         print(f"An error occurred: {str(e)}")
-        return frame, frame, None
+        return cv2.cvtColor(frame, cv2.COLOR_BGR2RGB), frame, None
 
 def process_video(input_video, selected_model, frame_index):
     try:
@@ -112,7 +109,7 @@ def process_video(input_video, selected_model, frame_index):
         if not ret:
             return None, None
 
-        _, processed_frame, graph = process_frame(frame, selected_model)
+        processed_frame, _, graph = process_frame(frame, selected_model)
         cap.release()
 
         return processed_frame, graph
@@ -162,6 +159,13 @@ with gr.Blocks() as app:
                 
                 def process_image(image, model):
                     frame_rgb, _, graph = process_frame(image, model)
+                    # Ensure the frame is in RGB format
+                    if len(frame_rgb.shape) == 3 and frame_rgb.shape[2] == 3:
+                        if frame_rgb.dtype != np.uint8:
+                            frame_rgb = (frame_rgb * 255).astype(np.uint8)
+                    else:
+                        # If the image is not in the correct format, convert it to RGB
+                        frame_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
                     return frame_rgb, graph
 
                 process_button_image.click(process_image, inputs=[image_input, model_dropdown_image], outputs=[output_image, output_image_graph])