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