Update app.py

app.py

@@ -1,105 +1,48 @@
-from ultralytics import YOLO
-import os
-
-# Define the correct path to config.yaml (in the root directory)
-config_path = './config.yaml'  # Adjust based on the actual path to your config.yaml
-
-# Load YOLO model
-model = YOLO("yolo11n.yaml")  # You can choose a different model type like yolo5n, yolo6n, etc.
-
-# Train the model
-results = model.train(data=config_path, epochs=1)
-
-# Define the save directory
-save_dir = './runs/detect/train/weights'
-
-# Create directory if it doesn't exist
-if not os.path.exists(save_dir):
-    os.makedirs(save_dir)
-
-# Save the model
-model.save(os.path.join(save_dir, 'best.pt'))
-
-# Print confirmation
-print("Model saved to:", os.path.join(save_dir, 'best.pt'))
-from ultralytics import YOLO
-import gradio as gr
+import torch
 import cv2
-import os
-import tempfile
-
-# Load the trained YOLO model
-model = YOLO("./runs/detect/train/weights/best.pt")  # Path to your trained model
+import numpy as np
 
-def process_video(video_path):
-    """
-    Process the input video using the YOLO model and save the output with bounding boxes.
-    Returns the path to the output video.
-    """
-    # Create a temporary file for the output video
-    output_path = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False).name
-    
-    # Open the input video
-    cap = cv2.VideoCapture(video_path)
-    if not cap.isOpened():
-        raise ValueError("Error opening video file")
+# Load your model (assuming it is a PyTorch model)
+model = torch.load('.data/model.pt')
+model.eval()
 
-    # Get video properties
-    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
-    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
-    fps = int(cap.get(cv2.CAP_PROP_FPS))
-    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
+# Open video file (input video)
+input_video = cv2.VideoCapture('input_video.mp4')
 
-    # Define the codec and create VideoWriter object
-    fourcc = cv2.VideoWriter_fourcc(*'mp4v')  # Use 'mp4v' for MP4 format
-    out = cv2.VideoWriter(output_path, fourcc, fps, (width, height))
+# Get the frame width, height, and frames per second (fps) from the input video
+frame_width = int(input_video.get(cv2.CAP_PROP_FRAME_WIDTH))
+frame_height = int(input_video.get(cv2.CAP_PROP_FRAME_HEIGHT))
+fps = input_video.get(cv2.CAP_PROP_FPS)
 
-    # Process each frame
-    while cap.isOpened():
-        ret, frame = cap.read()
-        if not ret:
-            break
+# Define the output video writer
+fourcc = cv2.VideoWriter_fourcc(*'mp4v')  # You can change this to any codec
+output_video = cv2.VideoWriter('output_video.mp4', fourcc, fps, (frame_width, frame_height))
 
-        # Perform YOLO inference on the frame
-        results = model(frame)
+while True:
+    # Read a frame from the input video
+    ret, frame = input_video.read()
+    if not ret:
+        break  # End of video
 
-        # Draw bounding boxes and labels on the frame
-        annotated_frame = results[0].plot()  # Ultralytics provides a plot method to draw boxes
+    # Preprocess the frame if necessary (depends on your model)
+    # For example, convert to tensor and normalize if required
+    frame_tensor = torch.tensor(frame).float().unsqueeze(0)  # Add batch dimension
 
-        # Write the annotated frame to the output video
-        out.write(annotated_frame)
+    # Pass the frame through the model
+    with torch.no_grad():
+        output = model(frame_tensor)  # Adjust based on your model's requirements
 
-    # Release resources
-    cap.release()
-    out.release()
-    cv2.destroyAllWindows()
+    # Postprocess the output if necessary (depends on your model's output format)
+    output_frame = output.squeeze(0).cpu().numpy()  # Remove batch dimension and convert to NumPy
 
-    return output_path
+    # Convert the model output to a valid image format (if necessary)
+    output_frame = np.uint8(output_frame)
 
-def gradio_interface(video):
-    """
-    Gradio interface function to handle video input and return the processed video.
-    """
-    if video is None:
-        return "Please upload a video file."
-    
-    try:
-        # Process the video and get the output path
-        output_video_path = process_video(video)
-        
-        # Return the output video for Gradio to display
-        return output_video_path
-    except Exception as e:
-        return f"Error processing video: {str(e)}"
+    # Write the frame to the output video
+    output_video.write(output_frame)
 
-# Create Gradio interface
-iface = gr.Interface(
-    fn=gradio_interface,
-    inputs=gr.Video(label="Upload Video"),
-    outputs=gr.Video(label="Processed Video with Detections"),
-    title="YOLOv11 Object Detection on Video",
-    description="Upload a video to run object detection using a trained YOLOv11 model."
-)
+# Release resources
+input_video.release()
+output_video.release()
 
-# Launch the Gradio interface
-iface.launch()
+cv2.destroyAllWindows()