Download player.py from tarto2/TurboVision: direct link, hf CLI and curl.
- Browser
- Download file 14.6 kB
-
https://huggingface.co/tarto2/TurboVision/resolve/main/player.py
- Command line
-
hf download hf://tarto2/TurboVision/player.py
-
curl -L -o player.py https://huggingface.co/tarto2/TurboVision/resolve/main/player.py
14.6 kB
| import cv2 | |
| import numpy as np | |
| from sklearn.cluster import KMeans | |
| import warnings | |
| import time | |
| import torch | |
| from torchvision.ops import batched_nms | |
| from numpy import ndarray | |
| # Suppress ALL runtime and sklearn warnings | |
| warnings.filterwarnings('ignore', category=RuntimeWarning) | |
| warnings.filterwarnings('ignore', category=FutureWarning) | |
| warnings.filterwarnings('ignore', category=UserWarning) | |
| # Suppress sklearn warnings specifically | |
| import logging | |
| logging.getLogger('sklearn').setLevel(logging.ERROR) | |
| def get_grass_color(img): | |
| # Convert image to HSV color space | |
| hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV) | |
| # Define range of green color in HSV | |
| lower_green = np.array([30, 40, 40]) | |
| upper_green = np.array([80, 255, 255]) | |
| # Threshold the HSV image to get only green colors | |
| mask = cv2.inRange(hsv, lower_green, upper_green) | |
| # Calculate the mean value of the pixels that are not masked | |
| masked_img = cv2.bitwise_and(img, img, mask=mask) | |
| grass_color = cv2.mean(img, mask=mask) | |
| return grass_color[:3] | |
| def get_players_boxes(frame, result): | |
| players_imgs = [] | |
| players_boxes = [] | |
| for (box, score, cls) in result: | |
| label = int(cls) | |
| if label == 0: | |
| x1, y1, x2, y2 = box.astype(int) | |
| player_img = frame[y1: y2, x1: x2] | |
| players_imgs.append(player_img) | |
| players_boxes.append([box, score, cls]) | |
| return players_imgs, players_boxes | |
| def get_kits_colors(players, grass_hsv=None, frame=None): | |
| kits_colors = [] | |
| if grass_hsv is None: | |
| grass_color = get_grass_color(frame) | |
| grass_hsv = cv2.cvtColor(np.uint8([[list(grass_color)]]), cv2.COLOR_BGR2HSV) | |
| for player_img in players: | |
| # Skip empty or invalid images | |
| if player_img is None or player_img.size == 0 or len(player_img.shape) != 3: | |
| continue | |
| # Convert image to HSV color space | |
| hsv = cv2.cvtColor(player_img, cv2.COLOR_BGR2HSV) | |
| # Define range of green color in HSV | |
| lower_green = np.array([grass_hsv[0, 0, 0] - 10, 40, 40]) | |
| upper_green = np.array([grass_hsv[0, 0, 0] + 10, 255, 255]) | |
| # Threshold the HSV image to get only green colors | |
| mask = cv2.inRange(hsv, lower_green, upper_green) | |
| # Bitwise-AND mask and original image | |
| mask = cv2.bitwise_not(mask) | |
| upper_mask = np.zeros(player_img.shape[:2], np.uint8) | |
| upper_mask[0:player_img.shape[0]//2, 0:player_img.shape[1]] = 255 | |
| mask = cv2.bitwise_and(mask, upper_mask) | |
| kit_color = np.array(cv2.mean(player_img, mask=mask)[:3]) | |
| kits_colors.append(kit_color) | |
| return kits_colors | |
| def get_kits_classifier(kits_colors): | |
| if len(kits_colors) == 0: | |
| return None | |
| if len(kits_colors) == 1: | |
| # Only one kit color, create a dummy classifier | |
| return None | |
| kits_kmeans = KMeans(n_clusters=2) | |
| kits_kmeans.fit(kits_colors) | |
| return kits_kmeans | |
| def classify_kits(kits_classifer, kits_colors): | |
| if kits_classifer is None or len(kits_colors) == 0: | |
| return np.array([0]) # Default to team 0 | |
| team = kits_classifer.predict(kits_colors) | |
| return team | |
| def get_left_team_label(players_boxes, kits_colors, kits_clf): | |
| left_team_label = 0 | |
| team_0 = [] | |
| team_1 = [] | |
| for i in range(len(players_boxes)): | |
| x1, y1, x2, y2 = players_boxes[i][0].astype(int) | |
| team = classify_kits(kits_clf, [kits_colors[i]]).item() | |
| if team == 0: | |
| team_0.append(np.array([x1])) | |
| else: | |
| team_1.append(np.array([x1])) | |
| team_0 = np.array(team_0) | |
| team_1 = np.array(team_1) | |
| # Safely calculate averages with fallback for empty arrays | |
| avg_team_0 = np.average(team_0) if len(team_0) > 0 else 0 | |
| avg_team_1 = np.average(team_1) if len(team_1) > 0 else 0 | |
| if avg_team_0 - avg_team_1 > 0: | |
| left_team_label = 1 | |
| return left_team_label | |
| def check_box_boundaries(boxes, img_height, img_width): | |
| """ | |
| Check if bounding boxes are within image boundaries and clip them if necessary. | |
| Args: | |
| boxes: numpy array of shape (N, 4) with [x1, y1, x2, y2] format | |
| img_height: height of the image | |
| img_width: width of the image | |
| Returns: | |
| valid_boxes: numpy array of valid boxes within boundaries | |
| valid_indices: indices of valid boxes | |
| """ | |
| x1, y1, x2, y2 = boxes[:, 0], boxes[:, 1], boxes[:, 2], boxes[:, 3] | |
| # Check if boxes are within boundaries | |
| valid_mask = (x1 >= 0) & (y1 >= 0) & (x2 < img_width) & (y2 < img_height) & (x1 < x2) & (y1 < y2) | |
| if not np.any(valid_mask): | |
| return np.array([]), np.array([]) | |
| valid_boxes = boxes[valid_mask] | |
| valid_indices = np.where(valid_mask)[0] | |
| # Clip boxes to image boundaries | |
| valid_boxes[:, 0] = np.clip(valid_boxes[:, 0], 0, img_width - 1) # x1 | |
| valid_boxes[:, 1] = np.clip(valid_boxes[:, 1], 0, img_height - 1) # y1 | |
| valid_boxes[:, 2] = np.clip(valid_boxes[:, 2], 0, img_width - 1) # x2 | |
| valid_boxes[:, 3] = np.clip(valid_boxes[:, 3], 0, img_height - 1) # y2 | |
| return valid_boxes, valid_indices | |
| def process_team_identification_batch(frames, results, kits_clf, left_team_label, grass_hsv): | |
| """ | |
| Process team identification and label formatting for batch results. | |
| Args: | |
| frames: list of frames | |
| results: list of detection results for each frame | |
| kits_clf: trained kit classifier | |
| left_team_label: label for left team | |
| grass_hsv: grass color in HSV format | |
| Returns: | |
| processed_results: list of processed results with team identification | |
| """ | |
| processed_results = [] | |
| for frame_idx, frame in enumerate(frames): | |
| frame_results = [] | |
| frame_detections = results[frame_idx] | |
| if not frame_detections: | |
| processed_results.append([]) | |
| continue | |
| # Extract player boxes and images | |
| players_imgs = [] | |
| players_boxes = [] | |
| player_indices = [] | |
| for idx, (box, score, cls) in enumerate(frame_detections): | |
| label = int(cls) | |
| if label == 0: # Player detection | |
| x1, y1, x2, y2 = box.astype(int) | |
| # Check boundaries | |
| if (x1 >= 0 and y1 >= 0 and x2 < frame.shape[1] and y2 < frame.shape[0] and x1 < x2 and y1 < y2): | |
| player_img = frame[y1:y2, x1:x2] | |
| if player_img.size > 0: # Ensure valid image | |
| players_imgs.append(player_img) | |
| players_boxes.append([box, score, cls]) | |
| player_indices.append(idx) | |
| # Initialize player team mapping | |
| player_team_map = {} | |
| # Process team identification if we have players | |
| if players_imgs and kits_clf is not None: | |
| kits_colors = get_kits_colors(players_imgs, grass_hsv) | |
| teams = classify_kits(kits_clf, kits_colors) | |
| # Create mapping from player index to team | |
| for i, team in enumerate(teams): | |
| player_team_map[player_indices[i]] = team.item() | |
| id = 0 | |
| # Process all detections with team identification | |
| for idx, (box, score, cls) in enumerate(frame_detections): | |
| label = int(cls) | |
| x1, y1, x2, y2 = box.astype(int) | |
| # Check boundaries | |
| valid_boxes, valid_indices = check_box_boundaries( | |
| np.array([[x1, y1, x2, y2]]), frame.shape[0], frame.shape[1] | |
| ) | |
| if len(valid_boxes) == 0: | |
| continue | |
| x1, y1, x2, y2 = valid_boxes[0].astype(int) | |
| # Apply team identification logic | |
| if label == 0: # Player | |
| if players_imgs and kits_clf is not None and idx in player_team_map: | |
| team = player_team_map[idx] | |
| if team == left_team_label: | |
| final_label = 6 # Player-L (Left team) | |
| else: | |
| final_label = 7 # Player-R (Right team) | |
| else: | |
| final_label = 6 # Default player label | |
| elif label == 1: # Goalkeeper | |
| final_label = 1 # GK | |
| elif label == 2: # Ball | |
| final_label = 0 # Ball | |
| elif label == 3 or label == 4: # Referee or other | |
| final_label = 3 # Referee | |
| else: | |
| continue | |
| # final_label = int(label) # Keep original label, ensure it's int | |
| frame_results.append({ | |
| "id": int(id), | |
| "bbox": [int(x1), int(y1), int(x2), int(y2)], | |
| "class_id": int(final_label), | |
| "conf": float(score) | |
| }) | |
| id = id + 1 | |
| processed_results.append(frame_results) | |
| return processed_results | |
| def convert_numpy_types(obj): | |
| """Convert numpy types to native Python types for JSON serialization.""" | |
| if isinstance(obj, np.integer): | |
| return int(obj) | |
| elif isinstance(obj, np.floating): | |
| return float(obj) | |
| elif isinstance(obj, np.ndarray): | |
| return obj.tolist() | |
| elif isinstance(obj, dict): | |
| return {key: convert_numpy_types(value) for key, value in obj.items()} | |
| elif isinstance(obj, list): | |
| return [convert_numpy_types(item) for item in obj] | |
| else: | |
| return obj | |
| def pre_process_img(frames, scale): | |
| imgs = np.stack([cv2.resize(frame, (int(scale), int(scale))) for frame in frames]) | |
| imgs = imgs.transpose(0, 3, 1, 2) | |
| imgs = imgs.astype(np.float32) / 255.0 # Normalize | |
| return imgs | |
| def post_process_output(outputs, x_scale, y_scale, conf_thresh=0.6, nms_thresh=0.75): | |
| B, C, N = outputs.shape | |
| outputs = torch.from_numpy(outputs) | |
| outputs = outputs.permute(0, 2, 1) | |
| boxes = outputs[..., :4] | |
| class_scores = 1 / (1 + torch.exp(-outputs[..., 4:])) | |
| conf, class_id = class_scores.max(dim=2) | |
| mask = conf > conf_thresh | |
| for i in range(class_id.shape[0]): # loop over batch | |
| # Find detections that are balls | |
| ball_idx = np.where(class_id[i] == 2)[0] | |
| if ball_idx.size > 0: | |
| # Pick the one with the highest confidence | |
| top = ball_idx[np.argmax(conf[i, ball_idx])] | |
| if conf[i, top] > 0.55: # apply confidence threshold | |
| mask[i, top] = True | |
| # ball_mask = (class_id == 2) & (conf > 0.51) | |
| # mask = mask | ball_mask | |
| batch_idx, pred_idx = mask.nonzero(as_tuple=True) | |
| if len(batch_idx) == 0: | |
| return [[] for _ in range(B)] | |
| boxes = boxes[batch_idx, pred_idx] | |
| conf = conf[batch_idx, pred_idx] | |
| class_id = class_id[batch_idx, pred_idx] | |
| x, y, w, h = boxes[:, 0], boxes[:, 1], boxes[:, 2], boxes[:, 3] | |
| x1 = (x - w / 2) * x_scale | |
| y1 = (y - h / 2) * y_scale | |
| x2 = (x + w / 2) * x_scale | |
| y2 = (y + h / 2) * y_scale | |
| boxes_xyxy = torch.stack([x1, y1, x2, y2], dim=1) | |
| max_coord = 1e4 | |
| offset = batch_idx.to(boxes_xyxy) * max_coord | |
| boxes_for_nms = boxes_xyxy + offset[:, None] | |
| keep = batched_nms(boxes_for_nms, conf, batch_idx, nms_thresh) | |
| boxes_final = boxes_xyxy[keep] | |
| conf_final = conf[keep] | |
| class_final = class_id[keep] | |
| batch_final = batch_idx[keep] | |
| results = [[] for _ in range(B)] | |
| for b in range(B): | |
| mask_b = batch_final == b | |
| if mask_b.sum() == 0: | |
| continue | |
| results[b] = list(zip(boxes_final[mask_b].numpy(), | |
| conf_final[mask_b].numpy(), | |
| class_final[mask_b].numpy())) | |
| return results | |
| def player_detection_result(frames: list[ndarray], batch_size, model, kits_clf=None, left_team_label=None, grass_hsv=None): | |
| start_time = time.time() | |
| # input_layer = model.input(0) | |
| # output_layer = model.output(0) | |
| height, width = frames[0].shape[:2] | |
| scale = 640.0 | |
| x_scale = width / scale | |
| y_scale = height / scale | |
| # infer_queue = AsyncInferQueue(model, len(frames)) | |
| infer_time = time.time() | |
| kits_clf = kits_clf | |
| left_team_label = left_team_label | |
| grass_hsv = grass_hsv | |
| results = [] | |
| for i in range(0, len(frames), batch_size): | |
| if i + batch_size > len(frames): | |
| batch_size = len(frames) - i | |
| batch_frames = frames[i:i + batch_size] | |
| imgs = pre_process_img(batch_frames, scale) | |
| input_name = model.get_inputs()[0].name | |
| outputs = model.run(None, {input_name: imgs})[0] | |
| raw_results = post_process_output(np.array(outputs), x_scale, y_scale) | |
| if kits_clf is None or left_team_label is None or grass_hsv is None: | |
| # Use first frame to initialize team classification | |
| first_frame = batch_frames[0] | |
| first_frame_results = raw_results[0] if raw_results else [] | |
| if first_frame_results: | |
| players_imgs, players_boxes = get_players_boxes(first_frame, first_frame_results) | |
| if players_imgs: | |
| grass_color = get_grass_color(first_frame) | |
| grass_hsv = cv2.cvtColor(np.uint8([[list(grass_color)]]), cv2.COLOR_BGR2HSV) | |
| kits_colors = get_kits_colors(players_imgs, grass_hsv) | |
| if kits_colors: # Only proceed if we have valid kit colors | |
| kits_clf = get_kits_classifier(kits_colors) | |
| if kits_clf is not None: | |
| left_team_label = int(get_left_team_label(players_boxes, kits_colors, kits_clf)) | |
| # Process team identification and boundary checking | |
| processed_results = process_team_identification_batch( | |
| batch_frames, raw_results, kits_clf, left_team_label, grass_hsv | |
| ) | |
| processed_results = convert_numpy_types(processed_results) | |
| results.extend(processed_results) | |
| # Return the same format as before for compatibility | |
| return results, kits_clf, left_team_label, grass_hsv |