change return type to model names
This commit is contained in:
		
							parent
							
								
									3224295610
								
							
						
					
					
						commit
						93f8974ba0
					
				
					 2 changed files with 6 additions and 1 deletions
				
			
		
							
								
								
									
										2
									
								
								.gitignore
									
										
									
									
										vendored
									
									
										Normal file
									
								
							
							
						
						
									
										2
									
								
								.gitignore
									
										
									
									
										vendored
									
									
										Normal file
									
								
							| 
						 | 
				
			
			@ -0,0 +1,2 @@
 | 
			
		|||
 | 
			
		||||
/__pycache__
 | 
			
		||||
							
								
								
									
										5
									
								
								app.py
									
										
									
									
									
								
							
							
						
						
									
										5
									
								
								app.py
									
										
									
									
									
								
							| 
						 | 
				
			
			@ -13,6 +13,9 @@ model = YOLO("yolov8n.pt")
 | 
			
		|||
if torch.cuda.is_available():
 | 
			
		||||
    model.to('cuda')
 | 
			
		||||
 | 
			
		||||
# Retrieve class names from the model
 | 
			
		||||
class_names = model.names
 | 
			
		||||
 | 
			
		||||
with open("config.json", "r") as f:
 | 
			
		||||
    config = json.load(f)
 | 
			
		||||
 | 
			
		||||
| 
						 | 
				
			
			@ -58,7 +61,7 @@ async def detect(websocket: WebSocket):
 | 
			
		|||
                for r in results:
 | 
			
		||||
                    for box in r.boxes:
 | 
			
		||||
                        boxes.append({
 | 
			
		||||
                            "class": int(box.cls[0]),
 | 
			
		||||
                            "class": class_names[int(box.cls[0])],  # Use class name
 | 
			
		||||
                            "confidence": float(box.conf[0]),
 | 
			
		||||
                        })
 | 
			
		||||
                # Broadcast to all subscribers of this URL
 | 
			
		||||
| 
						 | 
				
			
			
 | 
			
		|||
		Loading…
	
	Add table
		Add a link
		
	
		Reference in a new issue