Update app.py

app.py

@@ -16,10 +16,12 @@ def yolo_inference(input_file):
         results = model(img)
         annotated_img = results[0].plot()
 
-        # Save the annotated image
-        output_path = tempfile.NamedTemporaryFile(suffix=".png", delete=False).name
-        cv2.imwrite(output_path, annotated_img)
-        return output_path
+        # Display the annotated image in a window
+        cv2.imshow("YOLO Detection", annotated_img)
+        cv2.waitKey(0)
+        cv2.destroyAllWindows()
+
+        return input_file  # Return the original file for consistency (can be adjusted)
 
     elif input_file.endswith((".mp4", ".avi", ".mov")):
         # Process as a video
@@ -30,7 +32,8 @@ def yolo_inference(input_file):
         height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
         
         # Create a temporary output video path
-        output_video_path = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False).name
+        temp_dir = tempfile.mkdtemp()
+        output_video_path = os.path.join(temp_dir, "output.mp4")
         out = cv2.VideoWriter(output_video_path, fourcc, fps, (width, height))
         
         while cap.isOpened():
@@ -41,11 +44,20 @@ def yolo_inference(input_file):
             # Run YOLO on each frame
             results = model(frame)
             annotated_frame = results[0].plot()
+            
+            # Display the annotated frame in a window
+            cv2.imshow("YOLO Detection", annotated_frame)
+            if cv2.waitKey(1) & 0xFF == ord('q'):  # Press 'q' to quit early
+                break
+            
+            # Save the annotated frame to the video
             out.write(annotated_frame)
         
         cap.release()
         out.release()
-        return output_video_path
+        cv2.destroyAllWindows()
+
+        return input_file  # Return the original video file for consistency (can be adjusted)
 
     else:
         raise ValueError("Unsupported file format. Please upload an image or video.")
@@ -54,9 +66,9 @@ def yolo_inference(input_file):
 interface = gr.Interface(
     fn=yolo_inference,
     inputs=gr.File(label="Upload an Image or Video"),
-    outputs=gr.File(label="Processed Image/Video"),
+    outputs="text",  # Display a message about console output
     title="YOLO Object Detection",
-    description="Upload an image or video for object detection using YOLO."
+    description="Upload an image or video for object detection. The results are displayed on the console."
 )
 
 # Launch the app