Spaces:
Sleeping
Sleeping
Update app.py
Browse files
app.py
CHANGED
@@ -32,9 +32,14 @@ model.eval()
|
|
32 |
# model_path = "./last.pt"
|
33 |
# model = torch.jit.load(model_path, map_location=torch.device("cpu"))
|
34 |
# model.eval()
|
35 |
-
transform=transforms.Compose([
|
36 |
-
|
37 |
-
|
|
|
|
|
|
|
|
|
|
|
38 |
transforms.ToTensor()
|
39 |
])
|
40 |
# transform = transforms.Compose([
|
@@ -45,9 +50,22 @@ transform=transforms.Compose([
|
|
45 |
OBJECT_NAMES = ['enemies']
|
46 |
|
47 |
def detect_objects_in_image(image):
|
|
|
|
|
|
|
|
|
|
|
|
|
48 |
|
|
|
|
|
|
|
|
|
|
|
|
|
49 |
img_tensor = transform(image).unsqueeze(0)
|
50 |
-
|
|
|
51 |
|
52 |
with torch.no_grad():
|
53 |
pred = model(img_tensor)[0]
|
|
|
32 |
# model_path = "./last.pt"
|
33 |
# model = torch.jit.load(model_path, map_location=torch.device("cpu"))
|
34 |
# model.eval()
|
35 |
+
# transform=transforms.Compose([
|
36 |
+
# transforms.ToPILImage(),
|
37 |
+
# transforms.Resize((512,640)),
|
38 |
+
# transforms.ToTensor()
|
39 |
+
# ])
|
40 |
+
transform = transforms.Compose([
|
41 |
+
transforms.ToPILImage(), # Ensure input is a PIL image
|
42 |
+
transforms.Resize((512, 640)),
|
43 |
transforms.ToTensor()
|
44 |
])
|
45 |
# transform = transforms.Compose([
|
|
|
50 |
OBJECT_NAMES = ['enemies']
|
51 |
|
52 |
def detect_objects_in_image(image):
|
53 |
+
"""
|
54 |
+
Detect objects in the given image.
|
55 |
+
"""
|
56 |
+
# Ensure image is a PIL Image
|
57 |
+
if isinstance(image, torch.Tensor):
|
58 |
+
image = transforms.ToPILImage()(image) # Convert tensor to PIL image
|
59 |
|
60 |
+
if isinstance(image, Image.Image):
|
61 |
+
orig_w, orig_h = image.size # PIL image size returns (width, height)
|
62 |
+
else:
|
63 |
+
raise TypeError(f"Expected a PIL Image but got {type(image)}")
|
64 |
+
|
65 |
+
# Apply transformation
|
66 |
img_tensor = transform(image).unsqueeze(0)
|
67 |
+
|
68 |
+
|
69 |
|
70 |
with torch.no_grad():
|
71 |
pred = model(img_tensor)[0]
|