Update app.py

app.py

@@ -23,9 +23,29 @@ transform = transforms.Compose([
     transforms.ToTensor(),
 ])
 
+# COCO class names
+COCO_CLASSES = [
+    "person", "bicycle", "car", "motorcycle", "airplane", "bus", "train", "truck",
+    "boat", "traffic light", "fire hydrant", "stop sign", "parking meter", "bench",
+    "bird", "cat", "dog", "horse", "sheep", "cow", "elephant", "bear", "zebra",
+    "giraffe", "backpack", "umbrella", "handbag", "tie", "suitcase", "frisbee",
+    "skis", "snowboard", "sports ball", "kite", "baseball bat", "baseball glove",
+    "skateboard", "surfboard", "tennis racket", "bottle", "wine glass", "cup",
+    "fork", "knife", "spoon", "bowl", "banana", "apple", "sandwich", "orange",
+    "broccoli", "carrot", "hot dog", "pizza", "donut", "cake", "chair", "couch",
+    "potted plant", "bed", "dining table", "toilet", "TV", "laptop", "mouse",
+    "remote", "keyboard", "cell phone", "microwave", "oven", "toaster", "sink",
+    "refrigerator", "book", "clock", "vase", "scissors", "teddy bear", "hair drier",
+    "toothbrush"
+]
+
 # Prediction function
 def predict(image):
     try:
+        # Convert to RGB if the image has an alpha channel
+        if image.mode != 'RGB':
+            image = image.convert('RGB')
+
         # Transform the image
         image_tensor = transform(image).unsqueeze(0)  # Add batch dimension
         with torch.no_grad():
@@ -44,11 +64,12 @@ def predict(image):
         # Convert the input image to a NumPy array
         image_np = np.array(image)
 
-        # Draw boxes on the image
+        # Draw boxes and labels on the image
         for box, label in zip(boxes, labels):
             x1, y1, x2, y2 = box.astype(int)
             image_np = cv2.rectangle(image_np, (x1, y1), (x2, y2), (255, 0, 0), 2)
-            image_np = cv2.putText(image_np, str(label), (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 0, 0), 2)
+            class_name = COCO_CLASSES[label]  # Get the class name
+            image_np = cv2.putText(image_np, class_name, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 0, 0), 2)
 
         # Ensure the output is in the correct format
         return Image.fromarray(image_np.astype(np.uint8))