Update app.py
Browse files
app.py
CHANGED
@@ -44,9 +44,10 @@ def vision_ai_api(image, label):
|
|
44 |
}
|
45 |
|
46 |
def predict(image):
|
47 |
-
image = preprocess_image(image)
|
|
|
|
|
48 |
|
49 |
-
results = model(image, conf=0.85)
|
50 |
detected_classes = set()
|
51 |
labels = []
|
52 |
cropped_images = {}
|
@@ -57,37 +58,38 @@ def predict(image):
|
|
57 |
conf = box.conf[0]
|
58 |
cls = int(box.cls[0])
|
59 |
class_name = model.names[cls]
|
|
|
|
|
60 |
|
61 |
detected_classes.add(class_name)
|
62 |
labels.append(f"{class_name} {conf:.2f}")
|
63 |
|
64 |
-
#
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
65 |
cropped = image[y1:y2, x1:x2]
|
66 |
cropped_pil = Image.fromarray(cropped)
|
67 |
|
68 |
-
# Call
|
69 |
api_response = vision_ai_api(cropped_pil, class_name)
|
|
|
70 |
|
71 |
-
|
72 |
-
|
73 |
-
"image": cropped_pil,
|
74 |
-
"api_response": json.dumps(api_response, indent=4)
|
75 |
-
}
|
76 |
-
|
77 |
-
# Identify missing classes
|
78 |
-
possible_classes = {"front", "back"}
|
79 |
-
missing_classes = possible_classes - detected_classes
|
80 |
-
if missing_classes:
|
81 |
-
labels.append(f"Missing: {', '.join(missing_classes)}")
|
82 |
-
|
83 |
-
# Prepare Gradio outputs (separate front & back images and responses)
|
84 |
-
front_image = cropped_images.get("front", {}).get("image", None)
|
85 |
-
back_image = cropped_images.get("back", {}).get("image", None)
|
86 |
|
87 |
-
|
88 |
-
|
|
|
|
|
|
|
|
|
|
|
89 |
|
90 |
-
return front_image, front_response, back_image, back_response, labels
|
91 |
|
92 |
# Gradio Interface
|
93 |
iface = gr.Interface(
|
|
|
44 |
}
|
45 |
|
46 |
def predict(image):
|
47 |
+
image = preprocess_image(image)
|
48 |
+
|
49 |
+
results = model(image, conf=0.80)
|
50 |
|
|
|
51 |
detected_classes = set()
|
52 |
labels = []
|
53 |
cropped_images = {}
|
|
|
58 |
conf = box.conf[0]
|
59 |
cls = int(box.cls[0])
|
60 |
class_name = model.names[cls]
|
61 |
+
|
62 |
+
print(f"Detected: {class_name} ({conf:.2f}) at [{x1}, {y1}, {x2}, {y2}]")
|
63 |
|
64 |
detected_classes.add(class_name)
|
65 |
labels.append(f"{class_name} {conf:.2f}")
|
66 |
|
67 |
+
# Ensure bounding boxes are within the image
|
68 |
+
height, width = image.shape[:2]
|
69 |
+
x1, y1, x2, y2 = max(0, x1), max(0, y1), min(width, x2), min(height, y2)
|
70 |
+
|
71 |
+
if x1 >= x2 or y1 >= y2:
|
72 |
+
print("Invalid bounding box, skipping.")
|
73 |
+
continue
|
74 |
+
|
75 |
cropped = image[y1:y2, x1:x2]
|
76 |
cropped_pil = Image.fromarray(cropped)
|
77 |
|
78 |
+
# Call API
|
79 |
api_response = vision_ai_api(cropped_pil, class_name)
|
80 |
+
cropped_images[class_name] = {"image": cropped_pil, "api_response": json.dumps(api_response, indent=4)}
|
81 |
|
82 |
+
if not cropped_images:
|
83 |
+
return None, "No front detected", None, "No back detected", ["No valid detections"]
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
84 |
|
85 |
+
return (
|
86 |
+
cropped_images.get("front", {}).get("image", None),
|
87 |
+
cropped_images.get("front", {}).get("api_response", "{}"),
|
88 |
+
cropped_images.get("back", {}).get("image", None),
|
89 |
+
cropped_images.get("back", {}).get("api_response", "{}"),
|
90 |
+
labels
|
91 |
+
)
|
92 |
|
|
|
93 |
|
94 |
# Gradio Interface
|
95 |
iface = gr.Interface(
|