Update app.py

app.py

@@ -32,9 +32,14 @@ model.eval()
 # model_path = "./last.pt"
 # model = torch.jit.load(model_path, map_location=torch.device("cpu"))
 # model.eval()
-transform=transforms.Compose([
-    transforms.ToPILImage(),
-    transforms.Resize((512,640)),
+# transform=transforms.Compose([
+#     transforms.ToPILImage(),
+#     transforms.Resize((512,640)),
+#     transforms.ToTensor()
+# ])
+transform = transforms.Compose([
+    transforms.ToPILImage(),  # Ensure input is a PIL image
+    transforms.Resize((512, 640)),
     transforms.ToTensor()
 ])
 # transform = transforms.Compose([
@@ -45,9 +50,22 @@ transform=transforms.Compose([
 OBJECT_NAMES = ['enemies']
 
 def detect_objects_in_image(image):
+       """
+    Detect objects in the given image.
+    """
+    # Ensure image is a PIL Image
+    if isinstance(image, torch.Tensor):
+        image = transforms.ToPILImage()(image)  # Convert tensor to PIL image
     
+    if isinstance(image, Image.Image):
+        orig_w, orig_h = image.size  # PIL image size returns (width, height)
+    else:
+        raise TypeError(f"Expected a PIL Image but got {type(image)}")
+
+    # Apply transformation
     img_tensor = transform(image).unsqueeze(0)
-    orig_w, orig_h = image.size
+
+
 
     with torch.no_grad():
         pred = model(img_tensor)[0]