-
Notifications
You must be signed in to change notification settings - Fork 0
/
segtag_basicpose.py
149 lines (121 loc) · 5.49 KB
/
segtag_basicpose.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
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
import gradio as gr
from gradio import components as gc
import os
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'
import tensorflow as tf
from tensorflow.keras.applications import MobileNetV2
from tensorflow.keras.preprocessing import image
import numpy as np
import cv2
print("Starting script...")
# Set the GRADIO_SERVER_PORT environment variable
os.environ["GRADIO_SERVER_PORT"] = "7850"
# Explicitly define the input shape for MobileNetV2
input_shape = (224, 224, 3)
# Load the pre-trained MobileNetV2 model with the specified input shape
base_model = MobileNetV2(weights='imagenet', include_top=True, input_shape=input_shape)
shot_types = ["upper_body", "cowboy", "close_up", "portrait", "above", "full_body"]
def detect_damaged_eyes(image_path, blur_threshold):
image = cv2.imread(image_path)
gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
eye_cascade = cv2.CascadeClassifier(cv2.data.haarcascades + "haarcascade_eye.xml")
eyes = eye_cascade.detectMultiScale(gray, 1.3, 5)
value = 0 # Initialize the value variable here
for (ex, ey, ew, eh) in eyes:
roi_gray = gray[ey:ey+eh, ex:ex+ew]
value = cv2.Laplacian(roi_gray, cv2.CV_64F).var()
if value < blur_threshold: # Use the dynamic threshold
return True
return False
def classify_shot_type(features):
# Normalize the features to [0, 1] range
normalized_features = (features - np.min(features)) / (np.max(features) - np.min(features))
# Check the concentration of features
top_quarter = np.mean(normalized_features[:, :56])
upper_half = np.mean(normalized_features[:, :112])
lower_half = np.mean(normalized_features[:, 112:])
center = np.mean(normalized_features[112:160, 64:160])
if top_quarter > 0.6:
return "portrait"
elif center > 0.6:
return "close_up"
elif upper_half > lower_half + 0.1 and top_quarter < 0.6:
return "upper_body"
elif lower_half > upper_half + 0.1:
return "above"
elif upper_half > 0.5 and lower_half < 0.3:
return "cowboy_shot"
else:
return "full_body" # default to full body
def resize_and_crop(img, target_size=(224, 224)):
# Calculate aspect ratio
aspect = img.width / img.height
# Resize while maintaining aspect ratio
if aspect > 1:
# Landscape orientation - wide image
width = int(aspect * target_size[1])
img = img.resize((width, target_size[1]))
else:
# Portrait orientation - tall image
height = int(target_size[0] / aspect)
img = img.resize((target_size[0], height))
# Crop to the desired size
left_margin = (img.width - target_size[0]) / 2
bottom_margin = (img.height - target_size[1]) / 2
right_margin = left_margin + target_size[0]
top_margin = bottom_margin + target_size[1]
img = img.crop((left_margin, bottom_margin, right_margin, top_margin))
return img
def classify_image(img_path):
# Load the image
original_img = image.load_img(img_path)
# Create a copy of the original image for resizing and cropping
img = original_img.copy()
# Resize and crop the image for classification
img = resize_and_crop(img, target_size=(224, 224))
x = image.img_to_array(img)
x = np.expand_dims(x, axis=0)
x = tf.keras.applications.mobilenet_v2.preprocess_input(x)
# Extract features using MobileNetV2
features = base_model.predict(x)
# Classify the shot type based on the extracted features
shot_type = classify_shot_type(features)
return shot_type
def segregate_images(source_folder, upper_body, cowboy, close_up, portrait, above, full_body, blur_threshold):
# Create necessary directories within the source folder
folders = ["damaged_eyes"] + shot_types
for folder in folders:
if not os.path.exists(os.path.join(source_folder, folder)):
os.mkdir(os.path.join(source_folder, folder))
for root, dirs, files in os.walk(source_folder):
for file in files:
if file.endswith(('.png', '.jpg', '.jpeg')):
img_path = os.path.join(root, file)
if detect_damaged_eyes(img_path, blur_threshold): # Pass the dynamic threshold
os.rename(img_path, os.path.join(source_folder, "damaged_eyes", file))
else:
classification = classify_image(img_path)
os.rename(img_path, os.path.join(source_folder, classification, file))
return f"Images from {source_folder} have been processed."
def gradio_ui():
print("Initializing Gradio UI...")
interface = gr.Interface(
fn=segregate_images,
inputs=[
gc.Textbox(label="Angle Segregator: Provide the path to your source folder here. All images to be segregated need to be present in the source folder. Subfolders will not be scanned."),
gc.Checkbox(label="Upper Body"),
gc.Checkbox(label="Cowboy"),
gc.Checkbox(label="Close-Up"),
gc.Checkbox(label="Portrait"),
gc.Checkbox(label="Above"),
gc.Checkbox(label="Full Body"),
gc.Slider(minimum=0, maximum=100, default=50, label="Blur Threshold for Eyes") # Add the slider
],
outputs=gc.Textbox(),
live=False
)
print("Launching Gradio...")
interface.launch()
print("Executing main...")
if __name__ == "__main__":
gradio_ui()