Search code examples
pythonmachine-learningdeep-learningyolov5

YOLOv5 how to detect custom data


I wanted to do character detection in the game via yolov5 But so far I have just started to study it all How to train custom data, I know But I don't know how to tell yolo in the code so that it detects custom data in the screenshot that I make using mss, and numpy and I broadcast these screenshots using cv2.imshow thanks in advance

I tried different ways, but nothing worked, I looked in the official documentation. but without success


Solution

  • Using Custom Model

    import torch
    
    # Model
    # model = torch.hub.load('ultralytics/yolov5', 'yolov5s')  # or yolov5m, yolov5l, yolov5x, etc.
    model = torch.hub.load('ultralytics/yolov5', 'custom', 'path/to/best.pt')  # custom trained model
    
    # Images
    im = 'https://ultralytics.com/images/zidane.jpg'  # or file, Path, URL, PIL, OpenCV, numpy, list
    
    # Inference
    results = model(im)
    
    # Results
    results.print()  # or .show(), .save(), .crop(), .pandas(), etc.
    
    results.xyxy[0]  # im predictions (tensor)
    results.pandas().xyxy[0]  # im predictions (pandas)
    #      xmin    ymin    xmax   ymax  confidence  class    name
    # 0  749.50   43.50  1148.0  704.5    0.874023      0  person
    # 2  114.75  195.75  1095.0  708.0    0.624512      0  person
    # 3  986.00  304.00  1028.0  420.0    0.286865     27     tie