Update app.py

app.py

@@ -140,51 +140,75 @@ def match_and_identify(features, bbox):
         return identity, color
 
 def process_image(image):
-    # Prepare the image tensor
-    image_np = np.array(image)
-    input_tensor = np.expand_dims(image_np, axis=0)
-
-    # Run inference
-    detections = detect_objects(input_tensor)
-
-    # Extract output tensors and convert to numpy arrays
-    boxes = detections[0].numpy()[0]
-    scores = detections[1].numpy()[0]
-    classes = detections[2].numpy()[0]
-    num_detections = int(detections[3].numpy()[0])
-
-    # Filter detections for 'person' class
-    threshold = 0.3  # Adjust this threshold as needed
-    for i in range(num_detections):
-        class_id = int(classes[i])
-        score = scores[i]
-        box = boxes[i]
+    if image is None:
+        return None
 
-        if class_id == 1 and score > threshold:
-            h, w, _ = image.shape
-            ymin, xmin, ymax, xmax = box
-            left, right, top, bottom = int(xmin * w), int(xmax * w), int(ymin * h), int(ymax * h)
-            
-            # Extract person ROI
-            person_roi = image[top:bottom, left:right]
+    # Convert image to RGB if it's not
+    if len(image.shape) == 2:  # Grayscale
+        image = cv2.cvtColor(image, cv2.COLOR_GRAY2RGB)
+    elif image.shape[2] == 4:  # RGBA
+        image = cv2.cvtColor(image, cv2.COLOR_RGBA2RGB)
 
-            # Extract features
-            features = extract_features(person_roi)
+    # Ensure image is uint8
+    if image.dtype != np.uint8:
+        image = (image * 255).astype(np.uint8)
 
-            # Predict bbox using Kalman filter
-            predicted_bbox = np.array([xmin, ymin, xmax, ymax])
-
-            # Match and identify
-            identity, color = match_and_identify(features, predicted_bbox)
+    # Prepare the image tensor
+    image_np = np.array(image)
+    input_tensor = np.expand_dims(image_np, axis=0)
 
-            # Draw bounding box
-            left, top, right, bottom = int(predicted_bbox[0] * w), int(predicted_bbox[1] * h), int(predicted_bbox[2] * w), int(predicted_bbox[3] * h)
-            cv2.rectangle(image, (left, top), (right, bottom), color, 2)
-            cv2.putText(image, f'Person {identity}', (left, top - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)
+    try:
+        # Run inference
+        detections = detect_objects(input_tensor)
+
+        # Extract output tensors and convert to numpy arrays
+        boxes = detections[0].numpy()[0]
+        scores = detections[1].numpy()[0]
+        classes = detections[2].numpy()[0]
+        num_detections = int(detections[3].numpy()[0])
+
+        # Filter detections for 'person' class
+        threshold = 0.3  # Adjust this threshold as needed
+        for i in range(num_detections):
+            class_id = int(classes[i])
+            score = scores[i]
+            box = boxes[i]
+
+            if class_id == 1 and score > threshold:
+                h, w, _ = image.shape
+                ymin, xmin, ymax, xmax = box
+                left, right, top, bottom = int(xmin * w), int(xmax * w), int(ymin * h), int(ymax * h)
+                
+                # Extract person ROI
+                person_roi = image[top:bottom, left:right]
+
+                # Extract features
+                features = extract_features(person_roi)
+
+                # Predict bbox using Kalman filter
+                predicted_bbox = np.array([xmin, ymin, xmax, ymax])
+
+                # Match and identify
+                identity, color = match_and_identify(features, predicted_bbox)
+
+                # Draw bounding box
+                left, top, right, bottom = int(predicted_bbox[0] * w), int(predicted_bbox[1] * h), int(predicted_bbox[2] * w), int(predicted_bbox[3] * h)
+                cv2.rectangle(image, (left, top), (right, bottom), color, 2)
+                cv2.putText(image, f'Person {identity}', (left, top - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)
+
+    except Exception as e:
+        print(f"Error during processing: {str(e)}")
+        return image  # Return original image if there's an error
 
     return image
-
 def gradio_interface(input_image):
+    if input_image is None:
+        return None
+    
+    # Convert PIL Image to numpy array if necessary
+    if hasattr(input_image, 'convert'):
+        input_image = np.array(input_image.convert('RGB'))
+    
     # Process the input image
     output_image = process_image(input_image)
     return output_image