Spaces:
Sleeping
Sleeping
File size: 3,809 Bytes
8a03b37 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 | """
Model utilities for object detection
"""
import torch
import torchvision.transforms as transforms
# Exercise classes
CLASSES = [
"benchpress",
"deadlift",
"squat",
"leg_ext",
"pushup",
"shoulder_press"
]
# Image parameters
IMG_SIZE = 640 # Standard YOLO input size
# Allowed file extensions
ALLOWED_EXTENSIONS = {'png', 'jpg', 'jpeg', 'webp'}
def build_detection_model(model_type='yolov5', num_classes=6):
"""
Build object detection model
Args:
model_type: Type of model ('yolov5', 'yolov8', 'fasterrcnn')
num_classes: Number of exercise classes
Returns:
Detection model
"""
if model_type == 'yolov5':
# YOLOv5 model from ultralytics
model = torch.hub.load('ultralytics/yolov5', 'yolov5s', pretrained=True)
# Modify for custom classes if needed
model.model[-1].nc = num_classes
return model
elif model_type == 'yolov8':
from ultralytics import YOLO
model = YOLO('yolov8n.pt')
return model
elif model_type == 'fasterrcnn':
from torchvision.models.detection import fasterrcnn_resnet50_fpn
from torchvision.models.detection.faster_rcnn import FastRCNNPredictor
model = fasterrcnn_resnet50_fpn(pretrained=True)
in_features = model.roi_heads.box_predictor.cls_score.in_features
model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)
return model
else:
raise ValueError(f"Unknown model type: {model_type}")
def get_transforms(is_train=False):
"""
Get image transforms for detection
Args:
is_train: Whether transforms are for training
Returns:
Transform pipeline
"""
if is_train:
return transforms.Compose([
transforms.Resize((IMG_SIZE, IMG_SIZE)),
transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2),
transforms.RandomHorizontalFlip(p=0.5),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
])
else:
return transforms.Compose([
transforms.Resize((IMG_SIZE, IMG_SIZE)),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
])
def nms_boxes(boxes, scores, threshold=0.45):
"""
Apply Non-Maximum Suppression to filter overlapping boxes
Args:
boxes: List of bounding boxes [x1, y1, x2, y2]
scores: Confidence scores for each box
threshold: IOU threshold for suppression
Returns:
Indices of boxes to keep
"""
if len(boxes) == 0:
return []
boxes = torch.tensor(boxes, dtype=torch.float32)
scores = torch.tensor(scores, dtype=torch.float32)
keep = torch.ops.torchvision.nms(boxes, scores, threshold)
return keep.tolist()
def calculate_iou(box1, box2):
"""
Calculate Intersection over Union between two boxes
Args:
box1: [x1, y1, x2, y2]
box2: [x1, y1, x2, y2]
Returns:
IOU score
"""
x1 = max(box1[0], box2[0])
y1 = max(box1[1], box2[1])
x2 = min(box1[2], box2[2])
y2 = min(box1[3], box2[3])
intersection = max(0, x2 - x1) * max(0, y2 - y1)
box1_area = (box1[2] - box1[0]) * (box1[3] - box1[1])
box2_area = (box2[2] - box2[0]) * (box2[3] - box2[1])
union = box1_area + box2_area - intersection
return intersection / union if union > 0 else 0
|