Update app.py

app.py

@@ -27,8 +27,11 @@ def process_frame(frame, selected_model):
             model, preprocess = clip.load(selected_model, device=device)
             current_model_name = selected_model
         
-        # Convert frame to RGB
-        frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+        # Convert frame to RGB if it's not already
+        if len(frame.shape) == 3 and frame.shape[2] == 3:
+            frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+        else:
+            frame_rgb = frame
 
         # Detect faces
         boxes, _ = mtcnn.detect(Image.fromarray(frame_rgb))
@@ -41,7 +44,7 @@ def process_frame(frame, selected_model):
         # Process the largest face
         if largest_face is not None:
             x, y, w, h = map(int, largest_face)
-            cv2.rectangle(frame_rgb, (x, y), (w, h), (0, 0, 255), 2)
+            cv2.rectangle(frame_rgb, (x, y), (w, h), (0, 255, 0), 2)
             cropped_face = frame_rgb[y:h, x:w]
 
             # Convert the cropped face to a tensor
@@ -69,47 +72,43 @@ def process_frame(frame, selected_model):
             top_five_labels_probs = combined_labels_probs[:5]
 
             # Create a bar graph
-            fig, ax = plt.subplots(figsize=(10, 5), dpi=300)  # Increased figure size
-            plt.subplots_adjust(left=0.3)  # Adjust left margin
+            fig, ax = plt.subplots(figsize=(10, 5), dpi=300)
+            plt.subplots_adjust(left=0.3)
             
             labels, probs = zip(*top_five_labels_probs)
             bars = ax.barh(labels, probs, color=plt.cm.tab20.colors)
             ax.set_xlabel('Probability')
             ax.set_title('Top 5 Emotions')
-            ax.invert_yaxis()  # Invert y-axis to have the highest probability at the top
+            ax.invert_yaxis()
             
-            # Adjust x-axis labels to show only 3 decimal places
             ax.set_xticks(ax.get_xticks())
             ax.set_xticklabels([f'{x:.3f}' for x in ax.get_xticks()], rotation=0, ha='center')
             
-            # Ensure all labels are fully visible
             plt.tight_layout()
             
-            # Convert plot to image
             fig.canvas.draw()
             plot_img = np.frombuffer(fig.canvas.buffer_rgba(), dtype=np.uint8)
             plot_img = plot_img.reshape(fig.canvas.get_width_height()[::-1] + (4,))
             
             plt.close(fig)
         
-        return frame_rgb, frame_rgb, plot_img
+        return frame_rgb, plot_img
     except Exception as e:
         print(f"An error occurred: {str(e)}")
-        return cv2.cvtColor(frame, cv2.COLOR_BGR2RGB), frame, None
+        return frame_rgb, None
 
 def process_video(input_video, selected_model, frame_index):
     try:
         cap = cv2.VideoCapture(input_video)
         total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
         
-        # Set the frame position
         cap.set(cv2.CAP_PROP_POS_FRAMES, frame_index)
         
         ret, frame = cap.read()
         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
@@ -158,14 +157,7 @@ with gr.Blocks() as app:
                 output_image_graph = gr.Image(label="Results Graph")
                 
                 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)
+                    frame_rgb, graph = process_frame(image, model)
                     return frame_rgb, graph
 
                 process_button_image.click(process_image, inputs=[image_input, model_dropdown_image], outputs=[output_image, output_image_graph])