siam_rpn_tracker.py

import cv2
import numpy as np
import torch
import sys
sys.path.append('./Siamese-RPN-pytorch')  # Path to your cloned repo
from test import SiameseRPNTracker


class WebcamTracker:
    def __init__(self):
        self.tracker = None
        self.tracking = False
        self.click_point = None
        self.track_window = None
        self.frame = None
        self.window_name = "Siamese RPN Tracker"

        try:
            self.tracker = SiameseRPNTracker(model_path='path/to/your/model.pth')
            print("Tracker initialized successfully")
        except:
            print("Warning: Could not initialize tracker. Using dummy tracker for demonstration.")
            self.tracker = DummyTracker()

    def mouse_callback(self, event, x, y, flags, param):
        if event == cv2.EVENT_LBUTTONDOWN and not self.tracking:
            self.click_point = (x, y)
            self.start_tracking(x, y)

    def start_tracking(self, x, y):
        """Start tracking from the clicked point with 100x100 region"""
        if self.frame is None:
            return

        h, w = self.frame.shape[:2]

        bbox_size = 100
        x1 = max(0, x - bbox_size // 2)
        y1 = max(0, y - bbox_size // 2)
        x2 = min(w, x + bbox_size // 2)
        y2 = min(h, y + bbox_size // 2)

        bbox_w = x2 - x1
        bbox_h = y2 - y1

        if bbox_w < bbox_size:
            if x1 == 0:
                x2 = min(w, x1 + bbox_size)
            else:
                x1 = max(0, x2 - bbox_size)

        if bbox_h < bbox_size:
            if y1 == 0:
                y2 = min(h, y1 + bbox_size)
            else:
                y1 = max(0, y2 - bbox_size)

        bbox = [x1, y1, bbox_size, bbox_size]
        self.track_window = bbox

        small_frame = self.frame[y1:y1 + bbox_size, x1:x1 + bbox_size]

        try:
            self.tracker.init(self.frame, bbox)
            self.tracking = True
            print(f"Started tracking at ({x}, {y}) with 100x100 region")
            print(f"Bounding box: {bbox}")
            cv2.imshow("Tracking Region (100x100)", small_frame)
        except Exception as e:
            print(f"Error initializing tracker: {e}")
            self.tracking = False

    def run(self):
        """Main loop for webcam tracking"""
        cap = cv2.VideoCapture(0)

        if not cap.isOpened():
            print("Error: Could not open webcam")
            return

        cv2.namedWindow(self.window_name)
        cv2.setMouseCallback(self.window_name, self.mouse_callback)

        print("Instructions:")
        print("1. Click anywhere in the frame to select tracking point")
        print("2. Press 'q' to quit")
        print("3. Press 'r' to reset tracking")

        while True:
            ret, self.frame = cap.read()
            if not ret:
                print("Error: Could not read frame")
                break

            display_frame = self.frame.copy()

            if self.click_point and not self.tracking:
                x, y = self.click_point
                cv2.drawMarker(display_frame, (x, y), (0, 255, 0),
                               cv2.MARKER_CROSS, 20, 2)
                cv2.putText(display_frame, f"Click Point: ({x}, {y})",
                            (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 0), 2)

            if self.tracking:
                try:
                    bbox = self.tracker.update(self.frame)
                    x, y, w, h = map(int, bbox)

                    cv2.rectangle(display_frame, (x, y), (x + w, y + h),
                                  (0, 255, 0), 2)

                    center_x, center_y = x + w // 2, y + h // 2
                    cv2.circle(display_frame, (center_x, center_y),
                               3, (0, 0, 255), -1)

                    cv2.putText(display_frame,
                                f"Tracking: ({center_x}, {center_y})",
                                (10, 60), cv2.FONT_HERSHEY_SIMPLEX, 0.7,
                                (0, 255, 0), 2)

                    current_region = self.frame[y:y + h, x:x + w]
                    if current_region.size > 0:
                        if current_region.shape[:2] != (100, 100):
                            current_region = cv2.resize(current_region, (100, 100))
                        cv2.imshow("Current Tracked Region", current_region)

                except Exception as e:
                    print(f"Tracking error: {e}")
                    self.tracking = False

            cv2.putText(display_frame, "Click to select tracking point",
                        (10, display_frame.shape[0] - 10),
                        cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 255, 255), 1)

            cv2.imshow(self.window_name, display_frame)

            key = cv2.waitKey(1) & 0xFF
            if key == ord('q'):
                break
            elif key == ord('r'):
                self.tracking = False
                self.click_point = None
                self.track_window = None
                cv2.destroyWindow("Tracking Region (100x100)")
                cv2.destroyWindow("Current Tracked Region")
                print("Tracking reset")

        cap.release()
        cv2.destroyAllWindows()


class DummyTracker:
    """Fallback tracker for testing without the SiamRPN model"""
    def __init__(self):
        self.bbox = None

    def init(self, frame, bbox):
        self.bbox = bbox
        return True

    def update(self, frame):
        if self.bbox:
            x, y, w, h = self.bbox
            x += np.random.randint(-2, 3)
            y += np.random.randint(-2, 3)
            x = max(0, min(x, frame.shape[1] - w))
            y = max(0, min(y, frame.shape[0] - h))
            self.bbox = [x, y, w, h]
        return self.bbox


if __name__ == "__main__":
    tracker_app = WebcamTracker()
    tracker_app.run()