Skip to content

change in the predict.py #131

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 3 commits into
base: master
Choose a base branch
from
Open

change in the predict.py #131

Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
f1a2fef
Add files via upload
YasheshSavani Jan 19, 2018
639e6a6
Merge pull request #1 from Yashesh9274/Yashesh9274-patch-1-1
YasheshSavani Jan 19, 2018
d1f98e0
Add files via upload
YasheshSavani Jan 19, 2018
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
69 changes: 43 additions & 26 deletions predict.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
#! /usr/bin/env python

from __future__ import print_function
import argparse
import os
import cv2
Expand All @@ -9,9 +10,11 @@
from utils import draw_boxes
from frontend import YOLO
import json
import os
import os.path

os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"]=""
os.environ["CUDA_VISIBLE_DEVICES"]="0"

argparser = argparse.ArgumentParser(
description='Train and validate YOLO_v2 model on any dataset')
Expand Down Expand Up @@ -54,45 +57,59 @@ def _main_(args):
# Load trained weights
###############################

print weights_path
print (weights_path)
yolo.load_weights(weights_path)

###############################
# Predict bounding boxes
###############################

if image_path[-4:] == '.mp4':
video_out = image_path[:-4] + '_detected' + image_path[-4:]
##################
#A for loop for every item in [images] folder..

for file in os.listdir(image_path):

# Just Created another directory in output directory so input image and output image will be diffrentiated

if (os.path.isfile(os.path.join(image_path,file))):

video_reader = cv2.VideoCapture(image_path)
output_path = os.path.join(image_path,'output')
if not os.path.exists(output_path):
os.mkdir(output_path)

nb_frames = int(video_reader.get(cv2.CAP_PROP_FRAME_COUNT))
frame_h = int(video_reader.get(cv2.CAP_PROP_FRAME_HEIGHT))
frame_w = int(video_reader.get(cv2.CAP_PROP_FRAME_WIDTH))
if file[-4:] == '.mp4':
video_out = file[:-4] + '_detected' + file[-4:]

video_writer = cv2.VideoWriter(video_out,
cv2.VideoWriter_fourcc(*'MPEG'),
50.0,
(frame_w, frame_h))
video_reader = cv2.VideoCapture(os.path.join(image_path,file))

nb_frames = int(video_reader.get(cv2.CAP_PROP_FRAME_COUNT))
frame_h = int(video_reader.get(cv2.CAP_PROP_FRAME_HEIGHT))
frame_w = int(video_reader.get(cv2.CAP_PROP_FRAME_WIDTH))

for i in tqdm(range(nb_frames)):
_, image = video_reader.read()

boxes = yolo.predict(image)
image = draw_boxes(image, boxes, config['model']['labels'])
video_writer = cv2.VideoWriter(os.path.join(output_path,video_out),
cv2.VideoWriter_fourcc(*'MPEG'),
50.0,
(frame_w, frame_h))

video_writer.write(np.uint8(image))
for i in tqdm(range(nb_frames)):
_, image = video_reader.read()

boxes = yolo.predict(image)
image = draw_boxes(image, boxes, config['model']['labels'])

video_reader.release()
video_writer.release()
else:
image = cv2.imread(image_path)
boxes = yolo.predict(image)
image = draw_boxes(image, boxes, config['model']['labels'])
video_writer.write(np.uint8(image))

print len(boxes), 'boxes are found'
video_reader.release()
video_writer.release()
else:
image = cv2.imread(os.path.join(image_path,file)) ### joined the file name with the path
boxes = yolo.predict(image)
image = draw_boxes(image, boxes, config['model']['labels'])

cv2.imwrite(image_path[:-4] + '_detected' + image_path[-4:], image)
print (len(boxes), 'boxes are found')

new_file_name = file[:-4] + '_detected' + file[-4:] ####added new file name
cv2.imwrite(os.path.join(output_path, new_file_name),image)

if __name__ == '__main__':
args = argparser.parse_args()
Expand Down