Last active
September 22, 2020 19:04
-
-
Save sachinkale/57d6ababc8d9063646961af9659ac62e to your computer and use it in GitHub Desktop.
Extract Image from bounded box using YOLO
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
# import the necessary packages | |
from os import walk | |
import numpy as np | |
import argparse | |
import time | |
import cv2 | |
import os | |
import re | |
# construct the argument parse and parse the arguments | |
ap = argparse.ArgumentParser() | |
ap.add_argument("-y", "--yolo", required=True, | |
help="base path to YOLO directory") | |
ap.add_argument("-c", "--confidence", type=float, default=0.5, | |
help="minimum probability to filter weak detections") | |
ap.add_argument("-t", "--threshold", type=float, default=0.3, | |
help="threshold when applying non-maxima suppression") | |
args = vars(ap.parse_args()) | |
# load the COCO class labels our YOLO model was trained on | |
labelsPath = os.path.sep.join([args["yolo"], "coco.names"]) | |
LABELS = open(labelsPath).read().strip().split("\n") | |
# initialize a list of colors to represent each possible class label | |
np.random.seed(42) | |
COLORS = np.random.randint(0, 255, size=(len(LABELS), 3), | |
dtype="uint8") | |
# derive the paths to the YOLO weights and model configuration | |
weightsPath = os.path.sep.join([args["yolo"], "yolov3.weights"]) | |
configPath = os.path.sep.join([args["yolo"], "yolov3.cfg"]) | |
# load our YOLO object detector trained on COCO dataset (80 classes) | |
print("[INFO] loading YOLO from disk...") | |
net = cv2.dnn.readNetFromDarknet(configPath, weightsPath) | |
mypath = "/Users/xxx/work/scraping/test3" | |
cropoutputdir = "/Users/xxx/work/scraping/test3_cropped" | |
def cropimage(mypath, mydir, impath): | |
# load our input image and grab its spatial dimensions | |
imagepath = mypath + "/" + mydir + "/" + impath | |
croppedimagepath = cropoutputdir + "/" + mydir + "/" + impath | |
if not os.path.exists(cropoutputdir + "/" + mydir): | |
os.makedirs(cropoutputdir + "/" + mydir) | |
image = cv2.imread(imagepath) | |
(H, W) = image.shape[:2] | |
# determine only the *output* layer names that we need from YOLO | |
ln = net.getLayerNames() | |
ln = [ln[i[0] - 1] for i in net.getUnconnectedOutLayers()] | |
# construct a blob from the input image and then perform a forward | |
# pass of the YOLO object detector, giving us our bounding boxes and | |
# associated probabilities | |
blob = cv2.dnn.blobFromImage(image, 1 / 255.0, (416, 416), | |
swapRB=True, crop=False) | |
net.setInput(blob) | |
start = time.time() | |
layerOutputs = net.forward(ln) | |
end = time.time() | |
# show timing information on YOLO | |
print("[INFO] YOLO took {:.6f} seconds".format(end - start)) | |
# initialize our lists of detected bounding boxes, confidences, and | |
# class IDs, respectively | |
boxes = [] | |
confidences = [] | |
classIDs = [] | |
# loop over each of the layer outputs | |
for output in layerOutputs: | |
# loop over each of the detections | |
for detection in output: | |
# extract the class ID and confidence (i.e., probability) of | |
# the current object detection | |
scores = detection[5:] | |
classID = np.argmax(scores) | |
confidence = scores[classID] | |
# filter out weak predictions by ensuring the detected | |
# probability is greater than the minimum probability | |
if confidence > args["confidence"]: | |
# scale the bounding box coordinates back relative to the | |
# size of the image, keeping in mind that YOLO actually | |
# returns the center (x, y)-coordinates of the bounding | |
# box followed by the boxes' width and height | |
box = detection[0:4] * np.array([W, H, W, H]) | |
(centerX, centerY, width, height) = box.astype("int") | |
# use the center (x, y)-coordinates to derive the top and | |
# and left corner of the bounding box | |
x = int(centerX - (width / 2)) | |
y = int(centerY - (height / 2)) | |
# update our list of bounding box coordinates, confidences, | |
# and class IDs | |
boxes.append([x, y, int(width), int(height)]) | |
confidences.append(float(confidence)) | |
classIDs.append(classID) | |
# apply non-maxima suppression to suppress weak, overlapping bounding | |
# boxes | |
idxs = cv2.dnn.NMSBoxes(boxes, confidences, args["confidence"], | |
args["threshold"]) | |
# ensure at least one detection exists | |
if len(idxs) > 0: | |
# loop over the indexes we are keeping | |
for i in idxs.flatten(): | |
# extract the bounding box coordinates | |
(x, y) = (boxes[i][0], boxes[i][1]) | |
(w, h) = (boxes[i][2], boxes[i][3]) | |
# draw a bounding box rectangle and label on the image | |
color = [int(c) for c in COLORS[classIDs[i]]] | |
# cv2.rectangle(image, (x, y), (x + w, y + h), color, 2) | |
if(confidences[i] > 0.5 and LABELS[classIDs[i] == "Bird"]): | |
print("bird found") | |
print(mydir + "/" + impath) | |
crop = image[y:y+h, x:x+w] | |
if crop.size != 0: | |
cv2.imwrite(cropoutputdir + "/" + mydir + "/" + impath, crop) | |
break | |
# crop entire directory | |
for (dirpath, dirnames, filenames) in walk(mypath): | |
for mydir in dirnames: | |
print(mydir) | |
for(dirpath, dirnames, filenames) in walk(mypath + "/" + mydir): | |
for myfile in filenames: | |
if re.search('jpg', myfile): | |
cropimage(mypath, mydir, myfile) | |
break |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment