Skip to content

Commit

Permalink
견종 탐지 모델 설정
Browse files Browse the repository at this point in the history
  • Loading branch information
@sectionr0
sectionr0 committed Jun 1, 2022
1 parent 5950771 commit 92d904b
Show file tree
Hide file tree
Showing 2 changed files with 188 additions and 33 deletions.
200 changes: 171 additions & 29 deletions app.py
View file
Original file line number Diff line number Diff line change
@@ -1,10 +1,10 @@
import os
import os.path
import subprocess
import werkzeug

from flask import Flask, jsonify, send_file
import cv2
from flask import Flask, jsonify, send_file, request
from flask_restful import Resource, Api, reqparse
from werkzeug.sansio.multipart import MultipartEncoder

import yolov5.detect

Expand All @@ -28,10 +28,9 @@ def post(self):
file_object.save(dst=file_object.filename)
print(file_path)

breed = yolov5.detect.run(weights="yolov5/runs/train/dogtest5_yolov5m_results2/weights/best.pt",source=file_path)
breed = yolov5.detect.run(weights="yolov5/runs/train/dogtest5_yolov5m_results2/weights/best.pt", source=file_path)
safety = yolov5.detect.run(weights="yolov5/runs/train/dogsafety/best_safety.pt", source=file_path)
muzzle = yolov5.detect.run(weights="yolov5/runs/train/dogmuzzle/dog-muzzle-best.pt",
source=file_path)
muzzle = yolov5.detect.run(weights="yolov5/runs/train/dogmuzzle/dog-muzzle-best.pt", source=file_path)

# print("result", breed, safety, muzzle)

Expand All @@ -40,65 +39,208 @@ def post(self):
result_safety = False
result_muzzle = False


print("======result_breed======")
print(breed["result"])
count_breed_check1 = 0
count_breed_check2 = 0
count_breed_check3 = 0
count_breed_check4 = 0

for x in breed["result"]:
if "Rottweiler" in x:
result_muzzle = True
print("result_muzzle = True")
count_breed_check1 = count_breed_check1 + 1
# result_muzzle = True
# print("result_muzzle = True")
elif "tosa" in x:
result_muzzle = True
print("result_muzzle = True")
count_breed_check2 = count_breed_check2 + 1
# result_muzzle = True
# print("result_muzzle = True")
elif "dangerous_dog" in x:
result_muzzle = True
print("result_muzzle = True")
count_breed_check3 = count_breed_check3 + 1
# result_muzzle = True
# print("result_muzzle = True")
else:
result_muzzle = False
print("result_muzzle = False")
count_breed_check4 = count_breed_check4 + 1
# result_muzzle = False
# print("result_muzzle = False")
print("======result_breed222======")
print("count_breed_check1 ")
print(count_breed_check1)
print("count_breed_check2 ")
print(count_breed_check2)
print("count_breed_check3 ")
print(count_breed_check3)
print("count_breed_check4 ")
print(count_breed_check4)

max_count = max(count_breed_check1, count_breed_check2, count_breed_check3, count_breed_check4)

print("max_count ")
print(max_count)
print("======result_breed222======")

if max_count == count_breed_check1:
result_breed = True
print("result_breed = True")
elif max_count == count_breed_check2:
result_breed = True
print("result_breed = True")
elif max_count == count_breed_check3:
result_breed = True
print("result_breed = True")
else:
result_breed = False
print("result_breed = False")

print(breed["path"])
# img_breed = cv2.imread(breed["path"], cv2.IMREAD_COLOR)
# cv2.imshow('img_breed', img_breed)
# cv2.waitKey(0)
# cv2.destroyAllWindows()

print("======result_breed======")

# for abc in breed['result']:
# if "Dog" in abc:
# print(abc)
#
# for abc in breed['result']:
# print(abc)

print("======result_safety======")
print(safety["result"])
count_safety_check1 = 0
count_safety_check2 = 0

for x in safety["result"]:
if "leash-detect" in x:
result_safety = True
print("result_safety = True")
if "leash-detect-not" in x:
count_safety_check2 = count_safety_check2 + 1
else:
result_safety = False
print("result_safety = False")
count_safety_check1 = count_safety_check1 + 1

print("======result_safety222======")
print("count_safety_check1 ")
print(count_safety_check1)
print("count_safety_check2 ")
print(count_safety_check2)

max_count_safety = max(count_safety_check1, count_safety_check2)

print("max_count_safety ")
print(max_count_safety)

print("======result_safety222======")

if max_count_safety == count_safety_check1:
result_safety = True
print("result_safety = True")
else:
result_safety = False
print("result_safety = False")

print(safety["path"])
# img_safety = cv2.imread(safety["path"], cv2.IMREAD_COLOR)
# cv2.imshow('img_safety', img_safety)
# cv2.waitKey(0)
# cv2.destroyAllWindows()

print("======result_safety======")

print("======result_muzzle======")
print(muzzle["result"])
count_muzzle_check1 = 0
count_muzzle_check2 = 0

for x in muzzle["result"]:
if "muzzle" in x:
result_muzzle = True
print("result_muzzle = True")
if "no-muzzle" in x:
count_muzzle_check2 = count_muzzle_check2 + 1

else:
print("result_muzzle = False")
count_muzzle_check1 = count_muzzle_check1 + 1


print("======result_muzzle222======")
print("count_muzzle_check1 ")
print(count_muzzle_check1)
print("count_muzzle_check2 ")
print(count_muzzle_check2)

max_count_muzzle = max(count_muzzle_check1, count_muzzle_check2)
print("max_count_muzzle ")
print(max_count_muzzle)

print("======result_muzzle222======")

if max_count_muzzle == count_muzzle_check2:
result_muzzle = False
print("result_muzzle = True")
else:
result_muzzle = True
print("result_muzzle = False")

print(muzzle["path"])
print("======result_muzzle======")
# img_muzzle = cv2.imread(muzzle["path"], cv2.IMREAD_COLOR)
# cv2.imshow('img_muzzle', img_muzzle)
# cv2.waitKey(0)
# cv2.destroyAllWindows()

send_file(
print("======result_muzzle======")

)
# img = cv2.imread('lena.jpg', cv2.IMREAD_COLOR)
# return send_file(
# img_breed, mimetype='image/jpeg'
# # img_safety,
# # img_muzzle
# )

# return send_file(
# img_muzzle,
# as_attachment = True,
# attachment_filename = 'test.jpg',
# mimetype = 'image/jpeg'
# )

# m = MultipartEncoder(fields={'markdown': "> Block quotes are written like so in markdown.",
# 'toPersonEmail': '[email protected]',
# 'files': (img_muzzle, open(img_muzzle, 'rb'),
# 'image.png')})
#
# return (m.to_string(), {'Content-Type': m.content_type})

return jsonify({
'result': file_path,
'result_breed': result_breed,
'result_muzzle': result_muzzle,
'result_safety': result_safety
})
'result_safety': result_safety,
'result_breed_imgPath': breed['path'],
'result_muzzle_imgPath': muzzle['path'],
'result_safety_imgPath': safety['path']
})

# return send_file(breed["path"])


api.add_resource(Upload, '/upload')


@app.route('/download/breed')
def breed():
json = request.get_json()
return send_file(json['path'])


@app.route('/download/muzzle')
def muzzle():
json = request.get_json()
return send_file(json['path'])


@app.route('/download/safety')
def safety():
json = request.get_json()
return send_file(json['path'])


@app.route('/')
def hello_world(): # put application's code here
return 'Hello World!'
Expand Down
21 changes: 17 additions & 4 deletions yolov5/detect.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -83,9 +83,13 @@ def run(
if is_url and is_file:
source = check_file(source) # download

target = list()
result_txt123 = list()

# Directories
save_dir = increment_path(Path(project) / name, exist_ok=exist_ok) # increment run
(save_dir / 'labels' if save_txt else save_dir).mkdir(parents=True, exist_ok=True) # make dir
# (save_dir / 'images'if save_txt else save_dir).mkdir(parents=True, exist_ok=True)

# Load model
device = select_device(device)
Expand Down Expand Up @@ -140,10 +144,10 @@ def run(
p, im0, frame = path, im0s.copy(), getattr(dataset, 'frame', 0)


target = list()

p = Path(p) # to Path
save_path = str(save_dir / p.name) # im.jpg
# save_path = str(save_dir / 'images' / p.stem) + f'_{frame}' + '.jpg'

print(save_path, "save_Path")
txt_path = str(save_dir / 'labels' / p.stem) + ('' if dataset.mode == 'image' else f'_{frame}') # im.txt
Expand All @@ -160,16 +164,21 @@ def run(
n = (det[:, -1] == c).sum() # detections per class
s += f"{n} {names[int(c)]}{'s' * (n > 1)}, " # add to string


n_ = f"{n} {names[int(c)]} {'s' * (n > 1)}"
target.append(n_)
print(n_)

# Write results
for *xyxy, conf, cls in reversed(det):
if save_txt: # Write to file

xywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) / gn).view(-1).tolist() # normalized xywh
line = (cls, *xywh, conf) if save_conf else (cls, *xywh) # label format

result_txt123.append(('%g ' * len(line)).rstrip() % line)
# print("=========result_txt =========")
# print(('%g ' * len(line)).rstrip() % line)

with open(txt_path + '.txt', 'a') as f:
f.write(('%g ' * len(line)).rstrip() % line + '\n')

Expand Down Expand Up @@ -201,6 +210,7 @@ def run(
h = int(vid_cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
else: # stream
fps, w, h = 30, im0.shape[1], im0.shape[0]

save_path = str(Path(save_path).with_suffix('.mp4')) # force *.mp4 suffix on results videos
vid_writer[i] = cv2.VideoWriter(save_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, (w, h))
vid_writer[i].write(im0)
Expand All @@ -209,8 +219,11 @@ def run(
LOGGER.info(f'{s}Done. ({t3 - t2:.3f}s)')

print(f'{s}Done. ')
print(target)

# print(target)
# print("=========result_txt =========")
# print(result_txt123)
# print("=========result_txt =========")


# Print results
Expand All @@ -222,7 +235,7 @@ def run(
if update:
strip_optimizer(weights) # update model (to fix SourceChangeWarning)

resultYolo = {"result" : target, "path" : save_path}
resultYolo = {"result": target, "path": save_path, "result_txt123": result_txt123}
return resultYolo


Expand Down

0 comments on commit 92d904b

Please sign in to comment.