forked from simaiden/Clothing-Detection
-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathnew_image_demo.py
executable file
·67 lines (40 loc) · 1.56 KB
/
new_image_demo.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
import torch
import os
import cv2
from yolo.utils.utils import *
from predictors.YOLOv3 import YOLOv3Predictor
import glob
from tqdm import tqdm
import sys
from helpers.ImageLoader import load_images_from_folder
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
torch.cuda.empty_cache()
user_input = input("Please enter the name of the folder of images to crop: ")
#YOLO PARAMS
yolo_df2_params = { "model_def" : "yolov3-df2.cfg",
"weights_path" : "yolov3-df2_15000.weights",
"class_path":"df2.names",
"conf_thres" : 0.5,
"nms_thres" :0.6,
"img_size" : 416,
"device" : device}
#DATASET
dataset = 'df2'
yolo_params = yolo_df2_params
#Classes
classes = load_classes(yolo_params["class_path"])
detectron = YOLOv3Predictor(params=yolo_params)
path = user_input
images, filenames = load_images_from_folder(path)
detections = []
for i in range (len(images)):
detections.append(detectron.get_detections(images[i]))
for x1, y1, x2, y2, cls_conf, cls_pred in detections[i]:
if(classes[int(cls_pred)]=="short sleeve top" or classes[int(cls_pred)]=="long sleeve top"):
x1, y1, x2, y2 = int(x1), int(y1), int(x2), int(y2)
y1 = 0 if y1<0 else y1
new_img=images[i][y1:y2,x1-5:x2+20]
if(new_img.any()):
cv2.imwrite('Crops/'+'crop_'+filenames[i]+'.jpg', new_img)
img_id = path.split('/')[-1].split('.')[0]
print('Images Successfully Cropped and Saved')