syedfaisalabrar commited on
Commit
04dd650
·
verified ·
1 Parent(s): 14565fe

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +24 -22
app.py CHANGED
@@ -44,9 +44,10 @@ def vision_ai_api(image, label):
44
  }
45
 
46
  def predict(image):
47
- image = preprocess_image(image) # Apply preprocessing
 
 
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
- # Crop detected region
 
 
 
 
 
 
 
65
  cropped = image[y1:y2, x1:x2]
66
  cropped_pil = Image.fromarray(cropped)
67
 
68
- # Call Vision AI API separately for front & back
69
  api_response = vision_ai_api(cropped_pil, class_name)
 
70
 
71
- # Store cropped images & API response
72
- cropped_images[class_name] = {
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
- front_response = cropped_images.get("front", {}).get("api_response", "{}")
88
- back_response = cropped_images.get("back", {}).get("api_response", "{}")
 
 
 
 
 
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(