yolov5
              
                
                
            
          copied
				 3 changed files with 28 additions and 0 deletions
			
			
		@ -0,0 +1,5 @@ | 
				
			|||
from yolov5.yolov5 import Yolov5 | 
				
			|||
 | 
				
			|||
 | 
				
			|||
def yolov5(): | 
				
			|||
    return Yolov5() | 
				
			|||
@ -0,0 +1,5 @@ | 
				
			|||
matplotlib>=3.2.2 | 
				
			|||
opencv-python>=4.1.2 | 
				
			|||
torch>=1.7.0 | 
				
			|||
torchvision>=0.8.1 | 
				
			|||
seaborn | 
				
			|||
@ -0,0 +1,18 @@ | 
				
			|||
import torch | 
				
			|||
import numpy | 
				
			|||
from towhee.operator import NNOperator | 
				
			|||
 | 
				
			|||
@register(outputschema=['boxes', 'classes', 'scores']) | 
				
			|||
class Yolov5(NNOperator): | 
				
			|||
    def __init__(self, model_name): | 
				
			|||
        super().__init__() | 
				
			|||
        self._model = torch.hub.load("ultralytics/yolov5", model_name, pretrained=True) | 
				
			|||
 | 
				
			|||
    def __call__(self, img: numpy.ndarray): | 
				
			|||
        # Get object detection results with YOLOv5 model | 
				
			|||
        results = self._model(img) | 
				
			|||
         | 
				
			|||
        boxes = results.xyxy[0] | 
				
			|||
        classes = list(results.pandas().xyxy[0].name) | 
				
			|||
        scores = list(results.pandas().xyxy[0].confidence) | 
				
			|||
        return boxes, classes, scores | 
				
			|||
					Loading…
					
					
				
		Reference in new issue