Aishee
Aishee

Reputation: 125

How to crop the detected object after training using YOLO?

I am using YOLO for model training. I want to crop the detected object. For Darknet repository am using is: https://github.com/AlexeyAB/darknet/

For Detection and storing output coordinates in a text file am using this: !./darknet detector test data_for_colab/obj.data data_for_colab/yolov3-tiny-obj.cfg yolov3-tiny-obj_10000.weights -dont_show -ext_output < TEST.txt > result.txt Result.jpg

Upvotes: 2

Views: 8235

Answers (2)

Mahdi am
Mahdi am

Reputation: 1

In Yolov5 you can easily crop detected object:

    python detect.py --save-crop --source folderpath --weight yolov5s.pt

To crop specific class add --classes x (x is class index number)

   python detect.py --save-crop --classes 0 --source folderpath --weight yolov5s.pt

It crops only person detected an image (person has class index 0 in cocodataset)

Upvotes: 0

Ankur De
Ankur De

Reputation: 141

Considering in the TEST.txt file you have details as the sample image. You can use re module of python for text pattern detection, ie your "class_name".

Parsing the .txt file

import re
path='/content/darknet/result.txt'
myfile=open(path,'r')
lines=myfile.readlines()
pattern= "class_name"

for line in lines:
  if re.search(pattern,line):
    Cord_Raw=line
Cord=Cord_Raw.split("(")[1].split(")")[0].split("  ")

Now we will get the coordinates in a list.

Coordinate calculation

x_min=int(Cord[1])
x_max=x_min + int(Cord[5])
y_min=int(Cord[3])
y_max=y_min+ int(Cord[7])

Cropping from the actual image

import cv2
img = cv2.imread("Image.jpg")
crop_img = img[y_min:y_max, x_min:x_max]
cv2.imwrite("Object.jpg",crop_img)

Upvotes: 5

Related Questions