Welcome to DOFBOT SE repository

6.yolo garbage identification

Test according to the trained model, and the trained object name can be recognized.

1. Main code

Code path：/home/yahboom/Dofbot/6.AI_Visuall/7.Garbage identification.ipynb

The following code content needs to be executed according to the actual step. It cannot be run all at once. Running the last unit will directly exit the thread. The specific code can be viewed in the code path.

Import various libraries and model files

​x
#!/usr/bin/env python3
# coding: utf-8
import time
import torch
import rospy
import Arm_Lib
import cv2 as cv
import numpy as np
from time import sleep
from numpy import random

from utils.torch_utils import select_device

from models.experimental import attempt_load

from utils.general import non_max_suppression, scale_coords, xyxy2xywh, plot_one_box

model_path = '/home/yahboom/Dofbot/6.AI_Visual/model0.pt'
# Initialize
device = select_device()
# Load model
model = attempt_load(model_path, map_location=device)

Junk identification function

xxxxxxxxxx
def single_garbage_run(self, image):
        '''
        Execute the garbage identification function
        Execute garbage recognition function
        :param image: original image The original image
        :return: Recognized image, recognition information(name, pos) Recognized image, identification information (name, pos)
        '''
        while 1:
            self.frame = cv.resize(image, (640, 480))
            try: self.garbage_getName()
            except Exception: print("sqaure_pos empty")
            return self.frame

List of junk names:

xxxxxxxxxx
 def garbage_getName(self):
        name = "None"
        if self.status == 'waiting':
            self.frame, msg = self.garbage_identify.garbage_run(self.frame)
            for key, pos in msg.items(): name = key
            if name == "Zip_top_can":              (self.garbage_num, self.garbage_class) = ('00', '01')
            if name == "Old_school_bag":           (self.garbage_num, self.garbage_class) = ('01', '01')
            if name == "Newspaper":                (self.garbage_num, self.garbage_class) = ('02', '01')
            if name == "Book":                     (self.garbage_num, self.garbage_class) = ('03', '01')
            if name == "Toilet_paper":             (self.garbage_num, self.garbage_class) = ('04', '02')
            if name == "Peach_pit":                (self.garbage_num, self.garbage_class) = ('05', '02')
            if name == "Cigarette_butts":          (self.garbage_num, self.garbage_class) = ('06', '02')
            if name == "Disposable_chopsticks":    (self.garbage_num, self.garbage_class) = ('07', '02')
            if name == "Egg_shell":                (self.garbage_num, self.garbage_class) = ('08', '03')
            if name == "Apple_core":               (self.garbage_num, self.garbage_class) = ('09', '03')
            if name == "Watermelon_rind":          (self.garbage_num, self.garbage_class) = ('10', '03')
            if name == "Fish_bone":                (self.garbage_num, self.garbage_class) = ('11', '03')
            if name == "Expired_tablets":          (self.garbage_num, self.garbage_class) = ('12', '04')
            if name == "Expired_cosmetics":        (self.garbage_num, self.garbage_class) = ('13', '04')
            if name == "Used_batteries":           (self.garbage_num, self.garbage_class) = ('14', '04')
            if name == "Syringe":                  (self.garbage_num, self.garbage_class) = ('15', '04')
            if name == "None":                     (self.garbage_num, self.garbage_class) = ('None', 'None')

Main thread:

xxxxxxxxxx
def camera():
    # 打开摄像头 Open camera
    capture = cv.VideoCapture(0)
    # 当摄像头正常打开的情况下循环执行
    while capture.isOpened():
        try:
            _, img = capture.read()
            img = cv.resize(img, (640, 480))
            img = single_garbage.single_garbage_run(img)
            if model == 'Exit':
                cv.destroyAllWindows()
                capture.release()
                break
            imgbox.value = cv.imencode('.jpg', img)[1].tobytes()
        except KeyboardInterrupt:capture.release()

After the program block is run, you can see the following interface to identify the corresponding garbage.