500 lines
17 KiB
Python
500 lines
17 KiB
Python
"""
|
|
Runs on video and tracks objects across time, generating masks and polygons for multiple frames.
|
|
Takes one annotation of each asset and fetches the frames before and after it.
|
|
"""
|
|
import os
|
|
import cv2
|
|
import ast
|
|
import json
|
|
import torch
|
|
import numpy as np
|
|
|
|
from collections import deque
|
|
from scipy.optimize import linear_sum_assignment
|
|
from segment_anything import sam_model_registry, SamPredictor
|
|
|
|
# =========================================================
|
|
# CONFIG
|
|
# =========================================================
|
|
VIDEO_PATH = "/media/bs5/New Volume1/Brinda/Dataset/GPS_location/videos/2026_0330_094759_F.MP4"
|
|
JSON_PATH = "/media/bs5/New Volume1/Brinda/Dataset/GPS_location/output/2026_0330_094759_F.json"
|
|
OUTPUT_DIR = "/media/bs5/New Volume1/Brinda/Dataset/GPS_location/sam/"
|
|
|
|
FRAME_GAP = 5
|
|
NUM_FRAMES = 5
|
|
|
|
SAM_CHECKPOINT = "sam_vit_h_4b8939.pth"
|
|
MODEL_TYPE = "vit_h"
|
|
|
|
|
|
# =========================================================
|
|
# OUTPUT DIRS
|
|
# =========================================================
|
|
IMG_DIR = os.path.join(OUTPUT_DIR, "images")
|
|
MASK_DIR = os.path.join(OUTPUT_DIR, "masks")
|
|
EXTRACT_DIR = os.path.join(OUTPUT_DIR, "masked")
|
|
os.makedirs(OUTPUT_DIR, exist_ok=True)
|
|
os.makedirs(IMG_DIR, exist_ok=True)
|
|
os.makedirs(MASK_DIR, exist_ok=True)
|
|
os.makedirs(EXTRACT_DIR, exist_ok=True)
|
|
|
|
5
|
|
# ================== KALMAN FILTER ==================
|
|
|
|
class KalmanFilter:
|
|
"""Kalman filter for smoother track predictions"""
|
|
def __init__(self):
|
|
self.kf = cv2.KalmanFilter(8, 4)
|
|
|
|
# State transition matrix
|
|
self.kf.transitionMatrix = np.eye(8, dtype=np.float32)
|
|
self.kf.transitionMatrix[0, 4] = 1
|
|
self.kf.transitionMatrix[1, 5] = 1
|
|
self.kf.transitionMatrix[2, 6] = 1
|
|
self.kf.transitionMatrix[3, 7] = 1
|
|
|
|
# Measurement matrix
|
|
self.kf.measurementMatrix = np.zeros((4, 8), dtype=np.float32)
|
|
self.kf.measurementMatrix[0, 0] = 1
|
|
self.kf.measurementMatrix[1, 1] = 1
|
|
self.kf.measurementMatrix[2, 2] = 1
|
|
self.kf.measurementMatrix[3, 3] = 1
|
|
|
|
self.kf.processNoiseCov = np.eye(8, dtype=np.float32) * 0.01
|
|
self.kf.measurementNoiseCov = np.eye(4, dtype=np.float32) * 0.5
|
|
self.kf.errorCovPost = np.eye(8, dtype=np.float32)
|
|
|
|
def predict(self):
|
|
prediction = self.kf.predict()
|
|
return prediction[:4].flatten()
|
|
|
|
def update(self, measurement):
|
|
self.kf.correct(np.array(measurement, dtype=np.float32).reshape(4, 1))
|
|
|
|
|
|
# ================== TRACK CLASS ==================
|
|
|
|
class Track:
|
|
def __init__(self, tlwh, score, track_id, class_id=0, class_name='unknown'):
|
|
self.tlwh = np.asarray(tlwh, dtype=np.float32)
|
|
self.score = score
|
|
self.track_id = track_id
|
|
self.class_id = class_id
|
|
self.class_name = class_name
|
|
self.time_since_update = 0
|
|
self.hits = 1
|
|
self.age = 1
|
|
|
|
# Kalman filter for motion prediction
|
|
self.kf = KalmanFilter()
|
|
self.kf.kf.statePost = np.array([*tlwh, 0, 0, 0, 0], dtype=np.float32).reshape(8, 1)
|
|
|
|
# Track confidence history for stability
|
|
self.confidence_history = deque([score], maxlen=10)
|
|
self.avg_confidence = score
|
|
|
|
def predict(self):
|
|
"""Predict next position using Kalman filter"""
|
|
self.age += 1
|
|
self.time_since_update += 1
|
|
predicted = self.kf.predict()
|
|
|
|
if self.time_since_update < 5:
|
|
self.tlwh = predicted
|
|
|
|
def update(self, tlwh, score, class_id, class_name):
|
|
"""Update track with new detection"""
|
|
self.tlwh = np.asarray(tlwh, dtype=np.float32)
|
|
self.score = score
|
|
self.class_id = class_id
|
|
self.class_name = class_name
|
|
self.time_since_update = 0
|
|
self.hits += 1
|
|
|
|
self.confidence_history.append(score)
|
|
self.avg_confidence = np.mean(list(self.confidence_history))
|
|
|
|
self.kf.update(tlwh)
|
|
|
|
def is_stable(self):
|
|
"""Check if track is stable enough to display"""
|
|
return (self.hits >= 3 and
|
|
self.avg_confidence > 0.3 and
|
|
self.time_since_update == 0)
|
|
|
|
|
|
# ================== BYTETRACK TRACKER ==================
|
|
|
|
class ByteTrackEnhanced:
|
|
"""
|
|
Enhanced ByteTrack with per-class track ID sequencing
|
|
"""
|
|
def __init__(self, track_thresh=0.4, match_thresh=0.5, buffer_size=60,
|
|
high_thresh=0.5, low_thresh=0.1, class_names=None):
|
|
self.track_thresh = track_thresh
|
|
self.high_thresh = high_thresh
|
|
self.low_thresh = low_thresh
|
|
self.match_thresh = match_thresh
|
|
self.buffer_size = buffer_size
|
|
|
|
self.tracks = []
|
|
self.class_names = class_names if class_names else []
|
|
|
|
# Per-class track ID counters
|
|
self.next_id_per_class = {}
|
|
|
|
def get_next_id(self, class_id):
|
|
"""Get next track ID for specific class"""
|
|
if class_id not in self.next_id_per_class:
|
|
self.next_id_per_class[class_id] = 1
|
|
|
|
track_id = self.next_id_per_class[class_id]
|
|
self.next_id_per_class[class_id] += 1
|
|
return track_id
|
|
|
|
def get_class_name(self, class_id):
|
|
"""Get class name from class ID"""
|
|
if class_id < len(self.class_names):
|
|
return self.class_names[class_id]
|
|
return f'class_{class_id}'
|
|
|
|
def iou(self, a, b):
|
|
"""Calculate IoU between two boxes"""
|
|
ax, ay, aw, ah = a
|
|
bx, by, bw, bh = b
|
|
ax2, ay2 = ax + aw, ay + ah
|
|
bx2, by2 = bx + bw, by + bh
|
|
|
|
inter_x1 = max(ax, bx)
|
|
inter_y1 = max(ay, by)
|
|
inter_x2 = min(ax2, bx2)
|
|
inter_y2 = min(ay2, by2)
|
|
|
|
inter = max(0, inter_x2 - inter_x1) * max(0, inter_y2 - inter_y1)
|
|
union = aw * ah + bw * bh - inter
|
|
|
|
return inter / union if union > 0 else 0
|
|
|
|
def iou_distance(self, tracks, detections):
|
|
"""Compute IoU distance matrix"""
|
|
iou_mat = np.zeros((len(tracks), len(detections)), dtype=np.float32)
|
|
|
|
for i, t in enumerate(tracks):
|
|
for j, d in enumerate(detections):
|
|
iou = self.iou(t.tlwh, d[:4])
|
|
|
|
# Penalize class mismatch
|
|
if len(d) > 5 and t.class_id != int(d[5]):
|
|
iou *= 0.3
|
|
|
|
# Boost for active tracks
|
|
if t.time_since_update == 0:
|
|
iou *= 1.1
|
|
|
|
iou_mat[i, j] = iou
|
|
|
|
return iou_mat
|
|
|
|
def linear_assignment(self, cost_matrix, thresh):
|
|
"""Perform linear assignment with threshold"""
|
|
if cost_matrix.size == 0:
|
|
return np.empty((0, 2), dtype=int), tuple(range(cost_matrix.shape[0])), tuple(range(cost_matrix.shape[1]))
|
|
|
|
matches, unmatched_a, unmatched_b = [], [], []
|
|
row_ind, col_ind = linear_sum_assignment(-cost_matrix)
|
|
|
|
for r, c in zip(row_ind, col_ind):
|
|
if cost_matrix[r, c] < thresh:
|
|
unmatched_a.append(r)
|
|
unmatched_b.append(c)
|
|
else:
|
|
matches.append([r, c])
|
|
|
|
unmatched_a += list(set(range(cost_matrix.shape[0])) - set(row_ind))
|
|
unmatched_b += list(set(range(cost_matrix.shape[1])) - set(col_ind))
|
|
|
|
return np.array(matches), tuple(unmatched_a), tuple(unmatched_b)
|
|
|
|
def update(self, detections):
|
|
"""Two-stage association with stability fixes"""
|
|
# Predict all tracks
|
|
for t in self.tracks:
|
|
t.predict()
|
|
|
|
confirmed_tracks = [t for t in self.tracks if t.hits >= 2]
|
|
unconfirmed_tracks = [t for t in self.tracks if t.hits < 2]
|
|
|
|
# Separate high and low confidence detections
|
|
high_dets = [d for d in detections if d[4] >= self.high_thresh]
|
|
low_dets = [d for d in detections if self.low_thresh <= d[4] < self.high_thresh]
|
|
|
|
# First association
|
|
unmatched_tracks = []
|
|
unmatched_dets = []
|
|
|
|
if len(confirmed_tracks) > 0 and len(high_dets) > 0:
|
|
iou_mat = self.iou_distance(confirmed_tracks, high_dets)
|
|
matches, u_tracks, u_dets = self.linear_assignment(iou_mat, self.match_thresh)
|
|
|
|
for m in matches:
|
|
track_idx, det_idx = m[0], m[1]
|
|
class_id = int(high_dets[det_idx][5]) if len(high_dets[det_idx]) > 5 else 0
|
|
class_name = self.get_class_name(class_id)
|
|
confirmed_tracks[track_idx].update(high_dets[det_idx][:4],
|
|
high_dets[det_idx][4],
|
|
class_id,
|
|
class_name)
|
|
|
|
unmatched_tracks = [confirmed_tracks[i] for i in u_tracks]
|
|
unmatched_dets = [high_dets[i] for i in u_dets]
|
|
else:
|
|
unmatched_tracks = confirmed_tracks
|
|
unmatched_dets = high_dets
|
|
|
|
# Second association with low confidence
|
|
if len(unmatched_tracks) > 0 and len(low_dets) > 0:
|
|
iou_mat = self.iou_distance(unmatched_tracks, low_dets)
|
|
matches, u_tracks, _ = self.linear_assignment(iou_mat, self.match_thresh * 0.6)
|
|
|
|
for m in matches:
|
|
track_idx, det_idx = m[0], m[1]
|
|
class_id = int(low_dets[det_idx][5]) if len(low_dets[det_idx]) > 5 else 0
|
|
class_name = self.get_class_name(class_id)
|
|
unmatched_tracks[track_idx].update(low_dets[det_idx][:4],
|
|
low_dets[det_idx][4],
|
|
class_id,
|
|
class_name)
|
|
|
|
unmatched_tracks = [unmatched_tracks[i] for i in u_tracks]
|
|
|
|
# Deal with unconfirmed tracks
|
|
if len(unmatched_dets) > 0:
|
|
for t in unconfirmed_tracks:
|
|
iou_mat = self.iou_distance([t], unmatched_dets)
|
|
if iou_mat.size > 0:
|
|
best_match = np.argmax(iou_mat[0])
|
|
if iou_mat[0, best_match] > self.match_thresh * 0.5:
|
|
class_id = int(unmatched_dets[best_match][5]) if len(unmatched_dets[best_match]) > 5 else 0
|
|
class_name = self.get_class_name(class_id)
|
|
t.update(unmatched_dets[best_match][:4],
|
|
unmatched_dets[best_match][4],
|
|
class_id,
|
|
class_name)
|
|
unmatched_dets.pop(best_match)
|
|
|
|
# Create new tracks
|
|
for d in unmatched_dets:
|
|
if d[4] > self.track_thresh:
|
|
overlaps = False
|
|
for t in self.tracks:
|
|
if self.iou(t.tlwh, d[:4]) > 0.3:
|
|
overlaps = True
|
|
break
|
|
|
|
if not overlaps:
|
|
class_id = int(d[5]) if len(d) > 5 else 0
|
|
class_name = self.get_class_name(class_id)
|
|
track_id = self.get_next_id(class_id)
|
|
self.tracks.append(Track(d[:4], d[4], track_id, class_id, class_name))
|
|
|
|
# Remove old tracks
|
|
self.tracks = [t for t in self.tracks if t.time_since_update < self.buffer_size]
|
|
|
|
return self.tracks
|
|
|
|
|
|
|
|
# =========================================================
|
|
# CLEAN CLASS NAME
|
|
# =========================================================
|
|
def clean_name(name):
|
|
return name.replace("LEFT_", "").replace("RIGHT_", "").replace("_Start", "").replace("_Stop", "")
|
|
|
|
# =========================================================
|
|
# PARSE JSON
|
|
# =========================================================
|
|
def parse_json(path):
|
|
with open(path, 'r') as f:
|
|
content = f.read()
|
|
|
|
try:
|
|
data = json.loads(content) # proper JSON
|
|
except json.JSONDecodeError:
|
|
data = ast.literal_eval(content) # handles single quotes etc.
|
|
|
|
parsed = {}
|
|
|
|
for asset in data["Assets"]:
|
|
cls = clean_name(asset[0])
|
|
frame = int(asset[2])
|
|
|
|
x1, y1 = asset[3]
|
|
x2, y2 = asset[4]
|
|
|
|
if frame not in parsed:
|
|
parsed[frame] = []
|
|
|
|
parsed[frame].append({
|
|
"class": cls,
|
|
"bbox": [x1, y1, x2, y2],
|
|
"id": asset[1]
|
|
})
|
|
|
|
return parsed
|
|
|
|
# =========================================================
|
|
# SAM
|
|
# =========================================================
|
|
def load_sam():
|
|
device = "cuda" if torch.cuda.is_available() else "cpu"
|
|
sam = sam_model_registry[MODEL_TYPE](checkpoint=SAM_CHECKPOINT)
|
|
sam.to(device)
|
|
return SamPredictor(sam)
|
|
|
|
def get_mask(image, bbox, predictor):
|
|
masks, scores, _ = predictor.predict(
|
|
box=np.array(bbox),
|
|
multimask_output=True
|
|
)
|
|
return masks[np.argmax(scores)]
|
|
|
|
# =========================================================
|
|
# TEMPORAL FRAMES"""
|
|
Reads YOLO label--> convert to pixel format --> Feed the box to Sam ---> Get mask
|
|
"""
|
|
# =========================================================
|
|
def get_frames(center, total):
|
|
frames = []
|
|
for i in range(-NUM_FRAMES, NUM_FRAMES + 1):
|
|
f = center + i * FRAME_GAP
|
|
if 0 <= f < total:
|
|
frames.append(f)
|
|
return frames
|
|
|
|
# =========================================================
|
|
# MASK → POLYGON
|
|
# =========================================================
|
|
def mask_to_polygon(mask):
|
|
mask_uint8 = (mask * 255).astype(np.uint8)
|
|
contours, _ = cv2.findContours(mask_uint8, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
|
|
|
|
polygons = []
|
|
for cnt in contours:
|
|
if len(cnt) >= 3:
|
|
poly = cnt.squeeze().tolist()
|
|
polygons.append(poly)
|
|
|
|
return polygons
|
|
|
|
# =========================================================
|
|
# OVERLAY MASK
|
|
# =========================================================
|
|
def overlay_mask(image, mask):
|
|
overlay = image.copy()
|
|
color = np.array([255, 0, 0]) # RED
|
|
|
|
overlay[mask == 1] = overlay[mask == 1] * 0.5 + color * 0.5
|
|
return overlay.astype(np.uint8)
|
|
|
|
# =========================================================
|
|
# MAIN PIPELINE
|
|
# =========================================================
|
|
def process_video():
|
|
predictor = load_sam()
|
|
parsed = parse_json(JSON_PATH)
|
|
|
|
cap = cv2.VideoCapture(VIDEO_PATH)
|
|
total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
|
|
|
|
bbox_json = {}
|
|
poly_json = {}
|
|
|
|
for center_frame in parsed:
|
|
|
|
temporal_frames = get_frames(center_frame, total_frames)
|
|
|
|
for obj_idx, obj in enumerate(parsed[center_frame]):
|
|
cls = obj["class"]
|
|
init_bbox = obj["bbox"]
|
|
obj_id = obj["id"]
|
|
|
|
tracker = ByteTrackEnhanced(class_names=[cls])
|
|
|
|
x1, y1, x2, y2 = init_bbox
|
|
tlwh = [x1, y1, x2-x1, y2-y1]
|
|
|
|
init_det = np.array([[*tlwh, 0.9, 0]])
|
|
tracker.update(init_det)
|
|
|
|
for f in temporal_frames:
|
|
cap.set(cv2.CAP_PROP_POS_FRAMES, f)
|
|
ret, frame = cap.read()
|
|
if not ret:
|
|
continue
|
|
|
|
detections = init_det if f == center_frame else np.empty((0, 6))
|
|
tracks = tracker.update(detections)
|
|
|
|
if len(tracks) == 0:
|
|
continue
|
|
|
|
t = tracks[0]
|
|
x, y, w, h = t.tlwh
|
|
bbox = [int(x), int(y), int(x+w), int(y+h)]
|
|
|
|
image_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
|
|
predictor.set_image(image_rgb)
|
|
|
|
mask = get_mask(image_rgb, bbox, predictor)
|
|
mask_uint8 = (mask * 255).astype(np.uint8)
|
|
|
|
# SAVE MASK
|
|
cv2.imwrite(os.path.join(MASK_DIR, f"{f}.png"), mask_uint8)
|
|
|
|
# SAVE MASKED IMAGE
|
|
extracted = cv2.bitwise_and(image_rgb, image_rgb, mask=mask_uint8)
|
|
cv2.imwrite(os.path.join(EXTRACT_DIR, f"{f}.png"),
|
|
cv2.cvtColor(extracted, cv2.COLOR_RGB2BGR))
|
|
|
|
# OVERLAY IMAGE
|
|
overlay = overlay_mask(image_rgb, mask)
|
|
cv2.imwrite(os.path.join(IMG_DIR, f"{f}.png"),
|
|
cv2.cvtColor(overlay, cv2.COLOR_RGB2BGR))
|
|
|
|
# STORE BBOX JSON
|
|
if str(f) not in bbox_json:
|
|
bbox_json[str(f)] = {}
|
|
|
|
if cls not in bbox_json[str(f)]:
|
|
bbox_json[str(f)][cls] = []
|
|
|
|
bbox_json[str(f)][cls].append([str(obj_id), [bbox[0], bbox[1]], [bbox[2], bbox[3]]])
|
|
|
|
# STORE POLYGON JSON
|
|
polygons = mask_to_polygon(mask)
|
|
|
|
if str(f) not in poly_json:
|
|
poly_json[str(f)] = {}
|
|
|
|
if cls not in poly_json[str(f)]:
|
|
poly_json[str(f)][cls] = []
|
|
|
|
poly_json[str(f)][cls].append([str(obj_id), polygons])
|
|
|
|
print(f"[DONE] Frame {f}")
|
|
|
|
cap.release()
|
|
|
|
# SAVE JSON FILES
|
|
with open(os.path.join(OUTPUT_DIR, "bboxes.json"), "w") as f:
|
|
json.dump(bbox_json, f, indent=2)
|
|
|
|
with open(os.path.join(OUTPUT_DIR, "polygons.json"), "w") as f:
|
|
json.dump(poly_json, f, indent=2)
|
|
|
|
|
|
# =========================================================
|
|
# ENTRY
|
|
# =========================================================
|
|
if __name__ == "__main__":
|
|
process_video() |