Skip to content

Instantly share code, notes, and snippets.

@mengwangk
Created February 16, 2021 12:35
Show Gist options
  • Save mengwangk/1cfc8db3dfb30f94ec16301002ae6e4f to your computer and use it in GitHub Desktop.
Save mengwangk/1cfc8db3dfb30f94ec16301002ae6e4f to your computer and use it in GitHub Desktop.
YOLOV5 - Backend yolov5.py
import torch
import os
import logging
from dotenv import load_dotenv, find_dotenv
from PIL import Image
load_dotenv(find_dotenv())
yolo_model = os.getenv("YOLO_MODEL", "yolov5s")
logging.info(f"YOLO model - {yolo_model}")
model = torch.hub.load("ultralytics/yolov5", yolo_model, pretrained=True)
def yolov5(img):
"""Process a PIL image."""
# Inference
results = model(img)
detected_classes = []
names = results.names
if results.pred is not None:
pred = results.pred[0]
if pred is not None:
for c in pred[:, -1].unique():
n = (pred[:, -1] == c).sum()
detected_classes.append(f"{n} {names[int(c)]}{'s' * (n > 1)}")
logging.info(f"Detected classes: {detected_classes}")
rendered_imgs = results.render()
converted_img = Image.fromarray(rendered_imgs[0]).convert("RGB")
return detected_classes, converted_img
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment