Spaces:
Sleeping
Sleeping
| # app.py | |
| import gradio as gr | |
| import torch | |
| from torchvision import transforms, models | |
| import cv2 | |
| import numpy as np | |
| from PIL import Image | |
| from ultralytics import YOLO | |
| def load_models(): | |
| # Initialize YOLO | |
| yolo_model = YOLO('HockeyAI.pt') | |
| # Initialize SqueezeNet | |
| squeezenet_model = models.squeezenet1_1(weights=None) | |
| squeezenet_model.classifier[1] = torch.nn.Conv2d(512, 8, kernel_size=1) | |
| squeezenet_model.num_classes = 8 | |
| squeezenet_model.load_state_dict(torch.load('best_model_squezenet.pth', map_location=torch.device('cpu'))) | |
| squeezenet_model.eval() | |
| return yolo_model, squeezenet_model | |
| def process_image(input_image): | |
| if input_image is None: | |
| return None | |
| # Convert to numpy array if needed | |
| if isinstance(input_image, str): | |
| image = cv2.imread(input_image) | |
| image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) | |
| else: | |
| image = input_image.copy() | |
| # Initialize models | |
| yolo_model, squeezenet_model = load_models() | |
| # Class labels for direction | |
| class_labels = [ | |
| "Bottom", "Bottom_Left", "Bottom_Right", "Left", | |
| "Right", "Top", "Top_Left", "Top_Right" | |
| ] | |
| # Transform for SqueezeNet | |
| transform = transforms.Compose([ | |
| transforms.Resize((224, 224)), | |
| transforms.ToTensor(), | |
| transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) | |
| ]) | |
| # Run YOLO detection | |
| results = yolo_model(image) | |
| # Process each detection | |
| for box in results[0].boxes: | |
| xyxy = box.xyxy[0].cpu().numpy() | |
| conf = float(box.conf[0].cpu().numpy()) | |
| cls = int(box.cls[0].cpu().numpy()) | |
| # Process only if it's a player (class 4) and confidence is above threshold | |
| if cls == 4 and conf > 0.5: | |
| x1, y1, x2, y2 = map(int, xyxy) | |
| # Crop and process for direction classification | |
| if x2 > x1 and y2 > y1: | |
| cropped_array = image[y1:y2, x1:x2] | |
| if cropped_array.size > 0: | |
| cropped_image = Image.fromarray(cropped_array) | |
| # Predict direction | |
| image_tensor = transform(cropped_image).unsqueeze(0) | |
| with torch.no_grad(): | |
| output = squeezenet_model(image_tensor) | |
| direction_class = torch.argmax(output, dim=1).item() | |
| # Draw annotations | |
| cv2.rectangle(image, (x1, y1), (x2, y2), (0, 255, 0), 2) | |
| cv2.putText(image, f"{conf:.2f}", (x1, y1-10), | |
| cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2) | |
| # Draw direction arrow | |
| center_x, center_y = (x1 + x2) // 2, (y1 + y2) // 2 | |
| arrow_length = 80 # Increased from 50 to 80 | |
| direction = class_labels[direction_class] | |
| # Calculate arrow endpoint | |
| end_x, end_y = center_x, center_y | |
| if "Top" in direction: | |
| end_y = center_y - arrow_length | |
| elif "Bottom" in direction: | |
| end_y = center_y + arrow_length | |
| if "Left" in direction: | |
| end_x = center_x - arrow_length | |
| elif "Right" in direction: | |
| end_x = center_x + arrow_length | |
| cv2.arrowedLine(image, (center_x, center_y), (end_x, end_y), | |
| (255, 0, 0), 4, tipLength=0.4) | |
| return image | |
| # Create Gradio interface | |
| def gradio_interface(): | |
| with gr.Blocks() as iface: | |
| gr.Markdown("# Player Direction Detection") | |
| gr.Markdown("Upload an image to detect players and their movement directions") | |
| with gr.Row(): | |
| with gr.Column(): | |
| input_image = gr.Image(label="Input Image", type="numpy") | |
| with gr.Column(): | |
| output_image = gr.Image(label="Output Image") | |
| # Handle image processing | |
| input_image.change( | |
| fn=process_image, | |
| inputs=[input_image], | |
| outputs=[output_image] | |
| ) | |
| # Add example images if you have them | |
| gr.Examples( | |
| examples=["example-1.jpg", "example-2.jpg"], | |
| inputs=input_image, | |
| outputs=output_image, | |
| fn=process_image, | |
| cache_examples=True | |
| ) | |
| return iface | |
| if __name__ == "__main__": | |
| iface = gradio_interface() | |
| iface.launch() |