satellite_pc.py

import cv2
import numpy as np
import os
from flask import Flask, jsonify, Response
import time
from datetime import datetime
from ultralytics import YOLO

app = Flask(__name__)

# --- YOLOv8 Custom Model Setup ---
# Load your custom YOLOv8 model from 'best.pt'
model = YOLO("best.pt")

# Set your target class name as defined in your custom model.
TARGET_CLASS = "rover"  # Adjust this to match your model's label

# Define folders to store images.
CAPTURE_FOLDER = "captures"
ANNOTATED_FOLDER = "annotated"
os.makedirs(CAPTURE_FOLDER, exist_ok=True)
os.makedirs(ANNOTATED_FOLDER, exist_ok=True)

def get_timestamped_filename(folder, prefix, ext="jpg"):
    """Generate a filename with the current timestamp in the given folder."""
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    return os.path.join(folder, f"{prefix}_{timestamp}.{ext}")

# --- Global Camera Initialization ---
# Initialize the camera once at startup.
cap = cv2.VideoCapture(1)  # Adjust index if needed.
if not cap.isOpened():
    raise RuntimeError("Could not open camera on device index 1")
# Allow the camera to warm up.
time.sleep(0.5)

def capture_image():
    """
    Captures an image from the already-initialized camera, saves it with a timestamped filename,
    and returns the captured image.
    """
    ret, frame = cap.read()
    if not ret:
        raise RuntimeError("Failed to capture image from the camera")
    filename = get_timestamped_filename(CAPTURE_FOLDER, "capture")
    cv2.imwrite(filename, frame)
    print(f"Captured image saved as {filename}")
    return frame

def detect_target_yolo(image, conf_threshold=0.5, target_class=TARGET_CLASS):
    """
    Uses the YOLOv8 custom model to detect objects in the image and returns the bounding box and
    centroid of the detection matching the target class with the highest confidence.
    
    Args:
        image: BGR image from the camera.
        conf_threshold: Confidence threshold for detections.
        target_class: The class name of the target object.
    
    Returns:
        center (tuple): (center_x, center_y) of the detected target.
        box (tuple): (x, y, w, h) bounding box of the detected target.
        If no target is found, returns (None, None).
    """
    results = model(image, conf=conf_threshold)[0]  # Process the first image result

    target_box = None
    target_conf = -1
    target_center = None

    if results.boxes is not None:
        boxes = results.boxes.xyxy.cpu().numpy()  # [x1, y1, x2, y2]
        confidences = results.boxes.conf.cpu().numpy()
        class_ids = results.boxes.cls.cpu().numpy().astype(int)
        for i in range(len(boxes)):
            class_name = model.names[class_ids[i]]
            if class_name == target_class:
                if confidences[i] > target_conf:
                    target_conf = confidences[i]
                    x1, y1, x2, y2 = boxes[i]
                    target_box = (int(x1), int(y1), int(x2 - x1), int(y2 - y1))
                    target_center = (int((x1 + x2) / 2), int((y1 + y2) / 2))
    return target_center, target_box

@app.route('/find_rover', methods=['GET'])
def find_rover_endpoint():
    """
    Capture an image, use the YOLOv8 custom model to detect the target object,
    and return its centroid coordinates as JSON.
    """
    try:
        image = capture_image()
    except Exception as e:
        return jsonify({"error": str(e)}), 500

    center, box = detect_target_yolo(image)
    if center is None:
        return jsonify({"error": "Target object not found"}), 404
    else:
        return jsonify({"x": center[0], "y": center[1]})

@app.route('/render', methods=['GET'])
def render_endpoint():
    """
    Capture an image, use the YOLOv8 custom model to detect the target object,
    annotate the image with a bounding box and the centroid, save the annotated image with a timestamp,
    and return the annotated image as JPEG.
    """
    try:
        image = capture_image()
    except Exception as e:
        return jsonify({"error": str(e)}), 500

    center, box = detect_target_yolo(image)
    
    if center is not None and box is not None:
        x, y, w, h = box
        cv2.rectangle(image, (x, y), (x+w, y+h), (0, 255, 0), 2)
        cv2.circle(image, center, 5, (0, 255, 0), -1)
        cv2.putText(image, f"{TARGET_CLASS} ({center[0]}, {center[1]})", (x, y - 10),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)
    else:
        cv2.putText(image, "Target object not found", (10, 30),
                    cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 2)
    
    annotated_filename = get_timestamped_filename(ANNOTATED_FOLDER, "annotated_capture")
    cv2.imwrite(annotated_filename, image)
    print(f"Annotated image saved as {annotated_filename}")
    
    ret, jpeg = cv2.imencode(".jpg", image)
    if not ret:
        return jsonify({"error": "Could not encode image"}), 500
    return Response(jpeg.tobytes(), mimetype="image/jpeg")

if __name__ == '__main__':
    try:
        app.run(host="0.0.0.0", port=5000)
    finally:
        # Release the camera on shutdown.
        cap.release()